diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 54112071c6..57c1c3de1b 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -2424,9 +2424,14 @@ def _pid_alive(pid: int) -> bool: return True -def _pid_file_entries() -> "list[tuple[Path, int]]": - """(path, pid) per recorded server, including the legacy studio.pid.""" - entries = [] +def _pid_file_entries() -> "list[tuple[int, list[Path]]]": + """(pid, 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. + """ + by_pid: "dict[int, list[Path]]" = {} try: paths = sorted(STUDIO_HOME.glob(PID_FILE_GLOB)) + [_PID_FILE] except OSError: @@ -2441,11 +2446,11 @@ def _pid_file_entries() -> "list[tuple[Path, int]]": except (OSError, UnicodeDecodeError): continue if text.isdigit(): - entries.append((path, int(text))) + by_pid.setdefault(int(text), []).append(path) else: typer.echo(f"Ignoring invalid PID file {path.name}: {text}") path.unlink(missing_ok = True) - return entries + return list(by_pid.items()) def _signal_stop(pid: int) -> "str | None": @@ -2477,9 +2482,10 @@ def stop(): raise typer.Exit(0) signalled, failed = [], [] - for path, pid in entries: + for pid, paths in entries: if not _pid_alive(pid): - path.unlink(missing_ok = True) + for path in paths: + path.unlink(missing_ok = True) continue error = _signal_stop(pid) if error is not None: @@ -2487,7 +2493,7 @@ def stop(): 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)) + signalled.append((pid, paths)) if not signalled and not failed: typer.echo("No running Unsloth server found (cleaned up stale PID files).") @@ -2499,15 +2505,16 @@ def stop(): break time.sleep(0.5) for entry in list(pending): - path, pid = entry + pid, paths = entry if not _pid_alive(pid): - path.unlink(missing_ok = True) + for path in paths: + 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: + for pid, _paths in pending: typer.echo(f"Unsloth server (PID {pid}) is shutting down (may take a few seconds).") if failed: raise typer.Exit(1) diff --git a/unsloth_cli/tests/test_studio_stop.py b/unsloth_cli/tests/test_studio_stop.py index f3d10c7405..e7a6fe22db 100644 --- a/unsloth_cli/tests/test_studio_stop.py +++ b/unsloth_cli/tests/test_studio_stop.py @@ -89,6 +89,40 @@ def test_stop_does_not_leave_the_older_instance_running(monkeypatch, tmp_path): assert live == set() +def test_stop_signals_each_server_once(monkeypatch, tmp_path): + # A server writes its per-port file AND studio.pid. It stays alive while it + # shuts down gracefully, so a second SIGTERM would hit the SIG_DFL the first + # one installs and hard-kill it mid-cleanup. + 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) + killed = [] + monkeypatch.setattr(studio_mod.os, "kill", lambda pid, _sig: killed.append(pid)) + monkeypatch.setattr(sys, "platform", "linux") + _write_pid(tmp_path, "studio-8901-8550.pid", 8550) + _write_pid(tmp_path, "studio.pid", 8550) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert killed == [8550] + assert result.output.lower().count("sent shutdown signal") == 1 + + +def test_stop_removes_every_stale_file_for_one_pid(monkeypatch, tmp_path): + studio_mod, _live, killed = _install(monkeypatch, tmp_path, alive = set()) + _write_pid(tmp_path, "studio-8901-8550.pid", 8550) + _write_pid(tmp_path, "studio.pid", 8550) + + result = _run_stop(studio_mod) + + assert result.exit_code == 0, result.output + assert killed == [] + assert not list(tmp_path.glob("*.pid")) + + 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)