Stop every running Unsloth server, and refuse to start a second on a taken port
This commit is contained in:
Nilay Yadav
parent
52a9601032
commit
e5aec0aa05
4 changed files with 391 additions and 47 deletions
|
|
@ -744,7 +744,35 @@ def _find_free_port(
|
|||
|
||||
from utils.paths.storage_roots import studio_root as _studio_root
|
||||
|
||||
# Legacy single-instance file; still read so `stop` finds an older build's server.
|
||||
_PID_FILE = _studio_root() / "studio.pid"
|
||||
PID_FILE_GLOB = "studio-*.pid"
|
||||
|
||||
|
||||
def _pid_file_for_port(port: int) -> Path:
|
||||
return _studio_root() / f"studio-{port}.pid"
|
||||
|
||||
|
||||
def _blocker_is_own_studio(blocker: "tuple[int, str] | None") -> bool:
|
||||
"""True when the process holding the port is a server we recorded."""
|
||||
return bool(blocker) and blocker[0] in _recorded_studio_pids()
|
||||
|
||||
|
||||
def _recorded_studio_pids() -> "set[int]":
|
||||
"""PIDs recorded under this STUDIO_HOME."""
|
||||
pids: "set[int]" = set()
|
||||
try:
|
||||
paths = list(_studio_root().glob(PID_FILE_GLOB)) + [_PID_FILE]
|
||||
except OSError:
|
||||
return pids
|
||||
for path in paths:
|
||||
try:
|
||||
text = path.read_text(encoding = "utf-8").strip()
|
||||
except (OSError, UnicodeDecodeError):
|
||||
continue
|
||||
if text.isdigit():
|
||||
pids.add(int(text))
|
||||
return pids
|
||||
|
||||
# Direct backend launches bypass the CLI's env re-export; do it here for
|
||||
# real custom roots so unsloth-zoo's import-time LLAMA_CPP_DEFAULT_DIR
|
||||
|
|
@ -770,22 +798,30 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
|
|||
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
|
||||
|
||||
|
||||
def _write_pid_file():
|
||||
"""Write the current process PID to the studio PID file."""
|
||||
_OWN_PID_FILE: "Path | None" = None
|
||||
|
||||
|
||||
def _write_pid_file(port: int):
|
||||
"""Record this PID under its own port so `stop` can find every server."""
|
||||
global _OWN_PID_FILE
|
||||
path = _pid_file_for_port(port)
|
||||
try:
|
||||
_PID_FILE.parent.mkdir(parents = True, exist_ok = True)
|
||||
_PID_FILE.write_text(str(os.getpid()), encoding = "utf-8")
|
||||
path.parent.mkdir(parents = True, exist_ok = True)
|
||||
path.write_text(str(os.getpid()), encoding = "utf-8")
|
||||
except OSError:
|
||||
pass
|
||||
return
|
||||
_OWN_PID_FILE = path
|
||||
|
||||
|
||||
def _remove_pid_file():
|
||||
"""Remove the PID file if it belongs to this process."""
|
||||
if _OWN_PID_FILE is None:
|
||||
return
|
||||
try:
|
||||
if _PID_FILE.is_file():
|
||||
stored = _PID_FILE.read_text(encoding = "utf-8").strip()
|
||||
if _OWN_PID_FILE.is_file():
|
||||
stored = _OWN_PID_FILE.read_text(encoding = "utf-8").strip()
|
||||
if stored == str(os.getpid()):
|
||||
_PID_FILE.unlink(missing_ok = True)
|
||||
_OWN_PID_FILE.unlink(missing_ok = True)
|
||||
except (OSError, UnicodeDecodeError):
|
||||
pass
|
||||
|
||||
|
|
@ -1533,6 +1569,16 @@ def run_server(
|
|||
if not _is_port_free(host, port):
|
||||
original_port = port
|
||||
blocker = _get_pid_on_port(port)
|
||||
# Falling back past our own server is what creates the orphan.
|
||||
if _blocker_is_own_studio(blocker):
|
||||
print(
|
||||
f"Error: Unsloth Studio is already running on port {port} "
|
||||
f"(PID {blocker[0]}). Run `unsloth studio stop` first, or start this "
|
||||
"one on a different --port.",
|
||||
file = sys.stderr,
|
||||
flush = True,
|
||||
)
|
||||
sys.exit(1)
|
||||
port = _find_free_port(host, port + 1)
|
||||
if not silent:
|
||||
print("")
|
||||
|
|
@ -1731,7 +1777,7 @@ def run_server(
|
|||
(time.perf_counter() - boot_started) * 1000,
|
||||
)
|
||||
|
||||
_write_pid_file()
|
||||
_write_pid_file(port)
|
||||
import atexit
|
||||
|
||||
atexit.register(_remove_pid_file)
|
||||
|
|
|
|||
92
studio/backend/tests/test_studio_pid_files.py
Normal file
92
studio/backend/tests/test_studio_pid_files.py
Normal file
|
|
@ -0,0 +1,92 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""Per-port PID files, so `unsloth studio stop` can find every server.
|
||||
|
||||
Imports run.py directly, so run under the Unsloth venv.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
_BACKEND = Path(__file__).resolve().parents[1]
|
||||
if str(_BACKEND) not in sys.path:
|
||||
sys.path.insert(0, str(_BACKEND))
|
||||
|
||||
import run # noqa: E402
|
||||
|
||||
|
||||
@pytest.fixture(autouse = True)
|
||||
def isolated_root(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(run, "_studio_root", lambda: tmp_path)
|
||||
monkeypatch.setattr(run, "_PID_FILE", tmp_path / "studio.pid")
|
||||
monkeypatch.setattr(run, "_OWN_PID_FILE", None)
|
||||
yield
|
||||
|
||||
|
||||
def test_write_pid_file_is_per_port(tmp_path):
|
||||
run._write_pid_file(8901)
|
||||
|
||||
path = tmp_path / "studio-8901.pid"
|
||||
assert path.read_text(encoding = "utf-8") == str(os.getpid())
|
||||
|
||||
|
||||
def test_second_port_does_not_clobber_the_first(tmp_path):
|
||||
(tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8")
|
||||
|
||||
run._write_pid_file(8902)
|
||||
|
||||
assert (tmp_path / "studio-8901.pid").read_text(encoding = "utf-8") == "8550"
|
||||
assert (tmp_path / "studio-8902.pid").read_text(encoding = "utf-8") == str(os.getpid())
|
||||
|
||||
|
||||
def test_remove_pid_file_only_removes_our_own(tmp_path):
|
||||
run._write_pid_file(8901)
|
||||
(tmp_path / "studio-8902.pid").write_text("8600", encoding = "utf-8")
|
||||
|
||||
run._remove_pid_file()
|
||||
|
||||
assert not (tmp_path / "studio-8901.pid").exists()
|
||||
assert (tmp_path / "studio-8902.pid").exists()
|
||||
|
||||
|
||||
def test_remove_pid_file_leaves_a_reused_entry_alone(tmp_path):
|
||||
run._write_pid_file(8901)
|
||||
(tmp_path / "studio-8901.pid").write_text("999999", encoding = "utf-8")
|
||||
|
||||
run._remove_pid_file()
|
||||
|
||||
assert (tmp_path / "studio-8901.pid").read_text(encoding = "utf-8") == "999999"
|
||||
|
||||
|
||||
def test_recorded_studio_pids_reads_per_port_and_legacy_files(tmp_path):
|
||||
(tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8")
|
||||
(tmp_path / "studio-8902.pid").write_text("8600", encoding = "utf-8")
|
||||
(tmp_path / "studio.pid").write_text("4242", encoding = "utf-8")
|
||||
|
||||
assert run._recorded_studio_pids() == {8550, 8600, 4242}
|
||||
|
||||
|
||||
def test_recorded_studio_pids_ignores_corrupt_files(tmp_path):
|
||||
(tmp_path / "studio-8901.pid").write_text("not-a-pid", encoding = "utf-8")
|
||||
|
||||
assert run._recorded_studio_pids() == set()
|
||||
|
||||
|
||||
def test_own_studio_blocking_the_port_is_recognised(tmp_path):
|
||||
(tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8")
|
||||
|
||||
assert run._blocker_is_own_studio((8550, "python")) is True
|
||||
|
||||
|
||||
def test_a_foreign_blocker_still_falls_back(tmp_path):
|
||||
# jupyter-lab on 8888 must keep the fallback, not abort the launch.
|
||||
(tmp_path / "studio-8901.pid").write_text("8550", encoding = "utf-8")
|
||||
|
||||
assert run._blocker_is_own_studio((117, "jupyter-lab")) is False
|
||||
assert run._blocker_is_own_studio(None) is False
|
||||
|
|
@ -2394,6 +2394,7 @@ def run(
|
|||
# ── unsloth studio stop ───────────────────────────────────────────────
|
||||
|
||||
_PID_FILE = STUDIO_HOME / "studio.pid"
|
||||
PID_FILE_GLOB = "studio-*.pid"
|
||||
|
||||
|
||||
def _pid_alive(pid: int) -> bool:
|
||||
|
|
@ -2423,58 +2424,93 @@ def _pid_alive(pid: int) -> bool:
|
|||
return True
|
||||
|
||||
|
||||
@studio_app.command()
|
||||
def stop():
|
||||
"""Stop a running Unsloth Studio server.
|
||||
def _pid_file_entries() -> "list[tuple[Path, int]]":
|
||||
"""(path, pid) per recorded server, including the legacy studio.pid."""
|
||||
entries = []
|
||||
try:
|
||||
paths = sorted(STUDIO_HOME.glob(PID_FILE_GLOB)) + [_PID_FILE]
|
||||
except OSError:
|
||||
paths = [_PID_FILE]
|
||||
seen = set()
|
||||
for path in paths:
|
||||
if path in seen or not path.is_file():
|
||||
continue
|
||||
seen.add(path)
|
||||
try:
|
||||
text = path.read_text(encoding = "utf-8").strip()
|
||||
except (OSError, UnicodeDecodeError):
|
||||
continue
|
||||
if text.isdigit():
|
||||
entries.append((path, int(text)))
|
||||
else:
|
||||
typer.echo(f"Ignoring invalid PID file {path.name}: {text}")
|
||||
path.unlink(missing_ok = True)
|
||||
return entries
|
||||
|
||||
Reads the PID from ~/.unsloth/studio/studio.pid and sends SIGTERM
|
||||
(or TerminateProcess on Windows) to shut it down gracefully.
|
||||
"""
|
||||
|
||||
def _signal_stop(pid: int) -> "str | None":
|
||||
"""SIGTERM (or taskkill) the pid. Returns an error string, or None on success."""
|
||||
import signal as _signal
|
||||
|
||||
if not _PID_FILE.is_file():
|
||||
typer.echo("No running Unsloth server found (no PID file).")
|
||||
raise typer.Exit(0)
|
||||
|
||||
pid_text = _PID_FILE.read_text(encoding = "utf-8").strip()
|
||||
if not pid_text.isdigit():
|
||||
typer.echo(f"Invalid PID file contents: {pid_text}")
|
||||
_PID_FILE.unlink(missing_ok = True)
|
||||
raise typer.Exit(1)
|
||||
|
||||
pid = int(pid_text)
|
||||
|
||||
# Check if still alive (os.kill(pid, 0) is invalid on Windows -- see _pid_alive).
|
||||
if not _pid_alive(pid):
|
||||
typer.echo(f"Unsloth server (PID {pid}) is not running. Cleaning up stale PID file.")
|
||||
_PID_FILE.unlink(missing_ok = True)
|
||||
raise typer.Exit(0)
|
||||
|
||||
# Send SIGTERM (graceful shutdown) or TerminateProcess on Windows
|
||||
try:
|
||||
if sys.platform == "win32":
|
||||
# /T also stops llama-server children, which otherwise keep GPU and port.
|
||||
subprocess.run(["taskkill", "/PID", str(pid), "/T", "/F"], check = True)
|
||||
else:
|
||||
os.kill(pid, _signal.SIGTERM)
|
||||
typer.echo(f"Sent shutdown signal to Unsloth server (PID {pid}).")
|
||||
except ProcessLookupError:
|
||||
typer.echo(f"Unsloth server (PID {pid}) already exited.")
|
||||
_PID_FILE.unlink(missing_ok = True)
|
||||
raise typer.Exit(0)
|
||||
return None
|
||||
except Exception as e:
|
||||
typer.echo(f"Failed to stop Unsloth server (PID {pid}): {e}", err = True)
|
||||
raise typer.Exit(1)
|
||||
return str(e)
|
||||
return None
|
||||
|
||||
# Wait briefly for the process to exit and clean up.
|
||||
for _ in range(10):
|
||||
time.sleep(0.5)
|
||||
|
||||
@studio_app.command()
|
||||
def stop():
|
||||
"""Stop every running Unsloth Studio server for this STUDIO_HOME.
|
||||
|
||||
The port fallback can leave more than one running, so stop them all.
|
||||
"""
|
||||
entries = _pid_file_entries()
|
||||
if not entries:
|
||||
typer.echo("No running Unsloth server found (no PID file).")
|
||||
raise typer.Exit(0)
|
||||
|
||||
signalled, failed = [], []
|
||||
for path, pid in entries:
|
||||
if not _pid_alive(pid):
|
||||
_PID_FILE.unlink(missing_ok = True)
|
||||
typer.echo("Unsloth server stopped.")
|
||||
raise typer.Exit(0)
|
||||
path.unlink(missing_ok = True)
|
||||
continue
|
||||
error = _signal_stop(pid)
|
||||
if error is not None:
|
||||
failed.append((pid, error))
|
||||
typer.echo(f"Failed to stop Unsloth server (PID {pid}): {error}", err = True)
|
||||
continue
|
||||
typer.echo(f"Sent shutdown signal to Unsloth server (PID {pid}).")
|
||||
signalled.append((path, pid))
|
||||
|
||||
typer.echo("Unsloth server is shutting down (may take a few seconds).")
|
||||
if not signalled and not failed:
|
||||
typer.echo("No running Unsloth server found (cleaned up stale PID files).")
|
||||
raise typer.Exit(0)
|
||||
|
||||
pending = list(signalled)
|
||||
for _ in range(10):
|
||||
if not pending:
|
||||
break
|
||||
time.sleep(0.5)
|
||||
for entry in list(pending):
|
||||
path, pid = entry
|
||||
if not _pid_alive(pid):
|
||||
path.unlink(missing_ok = True)
|
||||
pending.remove(entry)
|
||||
|
||||
stopped = len(signalled) - len(pending)
|
||||
if stopped:
|
||||
typer.echo(f"Unsloth server{'s' if stopped > 1 else ''} stopped ({stopped}).")
|
||||
for _path, pid in pending:
|
||||
typer.echo(f"Unsloth server (PID {pid}) is shutting down (may take a few seconds).")
|
||||
if failed:
|
||||
raise typer.Exit(1)
|
||||
|
||||
|
||||
# ── unsloth studio setup / update ─────────────────────────────────────
|
||||
|
|
|
|||
170
unsloth_cli/tests/test_studio_stop.py
Normal file
170
unsloth_cli/tests/test_studio_stop.py
Normal file
|
|
@ -0,0 +1,170 @@
|
|||
# SPDX-License-Identifier: AGPL-3.0-only
|
||||
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
||||
|
||||
"""`unsloth studio stop` must stop every server it started.
|
||||
|
||||
With one PID file the second launch overwrote the first entry, so stop killed
|
||||
the newer server, claimed success, and left the older one serving.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from typer.testing import CliRunner
|
||||
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
if str(_REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(_REPO_ROOT))
|
||||
|
||||
|
||||
def _studio():
|
||||
from unsloth_cli.commands import studio as _studio_mod
|
||||
return _studio_mod
|
||||
|
||||
|
||||
def _install(monkeypatch, tmp_path, *, alive, killed = None):
|
||||
"""Point the CLI at tmp_path and fake process liveness."""
|
||||
studio_mod = _studio()
|
||||
monkeypatch.setattr(studio_mod, "STUDIO_HOME", tmp_path)
|
||||
monkeypatch.setattr(studio_mod, "_PID_FILE", tmp_path / "studio.pid")
|
||||
monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None)
|
||||
|
||||
live = set(alive)
|
||||
killed = killed if killed is not None else []
|
||||
|
||||
monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: pid in live)
|
||||
|
||||
def fake_kill(pid, _sig):
|
||||
killed.append(pid)
|
||||
live.discard(pid)
|
||||
|
||||
monkeypatch.setattr(studio_mod.os, "kill", fake_kill)
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
return studio_mod, live, killed
|
||||
|
||||
|
||||
def _write_pid(tmp_path, name, pid):
|
||||
(tmp_path / name).write_text(str(pid), encoding = "utf-8")
|
||||
|
||||
|
||||
def _run_stop(studio_mod):
|
||||
import typer as _typer
|
||||
|
||||
app = _typer.Typer()
|
||||
app.add_typer(studio_mod.studio_app, name = "studio")
|
||||
return CliRunner().invoke(app, ["studio", "stop"])
|
||||
|
||||
|
||||
def test_stop_kills_every_recorded_server(monkeypatch, tmp_path):
|
||||
studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = {8550, 8600})
|
||||
_write_pid(tmp_path, "studio-8901.pid", 8550)
|
||||
_write_pid(tmp_path, "studio-8902.pid", 8600)
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert sorted(killed) == [8550, 8600]
|
||||
assert not list(tmp_path.glob("studio-*.pid"))
|
||||
|
||||
|
||||
def test_stop_does_not_leave_the_older_instance_running(monkeypatch, tmp_path):
|
||||
# The reported symptom: stop claimed success while instance A kept serving.
|
||||
studio_mod, live, _killed = _install(monkeypatch, tmp_path, alive = {8550, 8600})
|
||||
_write_pid(tmp_path, "studio-8901.pid", 8550)
|
||||
_write_pid(tmp_path, "studio-8902.pid", 8600)
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert live == set()
|
||||
|
||||
|
||||
def test_stop_reads_the_legacy_single_pid_file(monkeypatch, tmp_path):
|
||||
studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = {4242})
|
||||
_write_pid(tmp_path, "studio.pid", 4242)
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert killed == [4242]
|
||||
assert not (tmp_path / "studio.pid").exists()
|
||||
|
||||
|
||||
def test_stop_reports_nothing_running_without_pid_files(monkeypatch, tmp_path):
|
||||
studio_mod, _live, _killed = _install(monkeypatch, tmp_path, alive = set())
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "no running unsloth server" in result.output.lower()
|
||||
|
||||
|
||||
def test_stop_cleans_stale_pid_files_without_claiming_a_stop(monkeypatch, tmp_path):
|
||||
studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = set())
|
||||
_write_pid(tmp_path, "studio-8901.pid", 8550)
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert killed == []
|
||||
assert not (tmp_path / "studio-8901.pid").exists()
|
||||
assert "stopped" not in result.output.lower()
|
||||
|
||||
|
||||
def test_stop_does_not_claim_a_stop_while_a_server_is_still_alive(monkeypatch, tmp_path):
|
||||
# SIGTERM delivered but it never exits: don't claim a stop, keep the file.
|
||||
studio_mod = _studio()
|
||||
monkeypatch.setattr(studio_mod, "STUDIO_HOME", 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.os, "kill", lambda pid, sig: None)
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
_write_pid(tmp_path, "studio-8901.pid", 8550)
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert "shutting down" in result.output.lower()
|
||||
assert "stopped" not in result.output.lower()
|
||||
assert (tmp_path / "studio-8901.pid").exists()
|
||||
|
||||
|
||||
def test_stop_continues_after_one_server_fails_to_stop(monkeypatch, tmp_path):
|
||||
studio_mod = _studio()
|
||||
monkeypatch.setattr(studio_mod, "STUDIO_HOME", tmp_path)
|
||||
monkeypatch.setattr(studio_mod, "_PID_FILE", tmp_path / "studio.pid")
|
||||
monkeypatch.setattr(studio_mod.time, "sleep", lambda _s: None)
|
||||
live = {8550, 8600}
|
||||
monkeypatch.setattr(studio_mod, "_pid_alive", lambda pid: pid in live)
|
||||
|
||||
def fake_kill(pid, _sig):
|
||||
if pid == 8550:
|
||||
raise PermissionError("not permitted")
|
||||
live.discard(pid)
|
||||
|
||||
monkeypatch.setattr(studio_mod.os, "kill", fake_kill)
|
||||
monkeypatch.setattr(sys, "platform", "linux")
|
||||
_write_pid(tmp_path, "studio-8901.pid", 8550)
|
||||
_write_pid(tmp_path, "studio-8902.pid", 8600)
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
combined = (result.output or "") + (getattr(result, "stderr", "") or "")
|
||||
assert result.exit_code == 1, combined
|
||||
assert 8600 not in live
|
||||
assert "8550" in combined
|
||||
|
||||
|
||||
def test_stop_discards_a_corrupt_pid_file(monkeypatch, tmp_path):
|
||||
studio_mod, _live, _killed = _install(monkeypatch, tmp_path, alive = set())
|
||||
(tmp_path / "studio-8901.pid").write_text("not-a-pid", encoding = "utf-8")
|
||||
|
||||
result = _run_stop(studio_mod)
|
||||
|
||||
assert result.exit_code == 0, result.output
|
||||
assert not (tmp_path / "studio-8901.pid").exists()
|
||||
Loading…
Add table
Add a link
Reference in a new issue