[staging CI] unslothai/unsloth#8170 - #221
danielhanchen wants to merge 62 commits into
Conversation
A Windows user could not update Studio until they killed a stray python by hand: the tool sandbox runs its payload under a shell wrapper, the kill path reaped only the wrapper, and `unsloth studio update` then refused to run because a process still held the managed environment. - taskkill /T on Windows, so a tool payload cannot outlive its wrapper - a console-close handler, since CTRL_CLOSE_EVENT never becomes a Python signal and the graceful shutdown was skipped entirely when the window was closed - the desktop updater drains the app job before standing down crash cleanup, and re-arms it when the install never happens - children are recorded on disk and swept at the next startup, which is the only reaper macOS has after a crash or a force quit - the job status is logged instead of failing silently Also fixes a liveness probe that used os.kill(pid, 0); on Windows that is TerminateProcess, so it killed the process it was asking about.
for more information, see https://pre-commit.ci
…d per owner, drop the job drain - the console handler ran _signal_handler on the thread Windows creates for the event, where signal.signal raises, so closing the window did no cleanup at all - bound that work to the ~5s Windows allows before it kills the process - one record file per owner pid: two Studios can share a home, and a single file let the second erase the first's children - add a Windows process identity (creation time) and refuse to signal a pid that cannot be verified - drop the whole-job drain: it would also terminate the WebView2 hosts, and cleanup_child_processes already taskkills the backend tree
A record that is not an object, or whose children are not dicts, raised out of the startup sweep and would have stopped Studio from starting. Pair the Linux start time with the command name as well, since start time alone has 10ms granularity.
for more information, see https://pre-commit.ci
Ctrl+C on Windows raised UnboundLocalError inside the console callback, where the BOOL result is then undefined, so the event could be reported as handled and Studio would not stop. The updater no longer re-arms kill-on-close before relaunching, which would have made the old process kill the replacement it just started, and it resets the exit-cleanup guard so a retry after a failed installer still reaps the backend. The RAG embedder and cloudflared are recorded like the other sidecars, a llama-server that survived a failed kill stays recorded, the Windows kill path checks the captured creation time before taskkill, and record writes are serialised.
for more information, see https://pre-commit.ci
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c6a8562842
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Deliberately NOT re-arming kill-on-close before the restart: relaunch() | ||
| // starts the replacement as a child, so it inherits this job, and re-arming | ||
| // would make this process kill it on the way out. |
There was a problem hiding this comment.
When the app was launched with --hidden, mark_in_app_relaunch can reject because its config directory or marker cannot be written. This call is outside the relaunch try, after the updater's on_before_exit hook has already consumed the termination guard and disabled Windows kill-on-close, so the outer recovery path restarts the backend without invoking resume_desktop_update_cleanup. The app then continues with no crash reaper, and a later force-exit can orphan that backend (while an update retry also skips its pre-installer cleanup); include marker/import failures in the same re-arm path.
Useful? React with 👍 / 👎.
The delayed Windows kill skips a captured pid it cannot verify, since the job object still takes the tree at exit. An owner whose identity cannot be read counts as live rather than gone, so a momentary ps failure no longer costs a running Studio its sidecars, and that lookup pins TZ so a timezone change does not read as a different process. A child that outlived terminate_all keeps its record instead of losing the only handle on it, with zombies told apart from survivors. The whole relaunch handoff is inside the recovery scope, so any path that leaves this process running re-arms cleanup.
for more information, see https://pre-commit.ci
c6a8562 to
6c1efba
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6c1efba31b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for path in sorted(directory.glob("*.json")): | ||
| killed.extend(_reap_one_record(path, timeout)) |
There was a problem hiding this comment.
On macOS, nested processes create dependent records: the backend records a multiprocessing worker, while that worker records its llama-server. If the worker's filename sorts before the backend's (for example across a PID digit boundary or after PID wrap), this single pass first skips the worker record because its owner is still alive, then kills the worker while processing the backend record, and never revisits the skipped record. Because macOS has no parent-death signal, the llama-server can remain alive until another Studio launch; sweep until no newly-dead owners are found rather than relying on record ordering.
Useful? React with 👍 / 👎.
| } catch (installError) { | ||
| // Failed or cancelled: we keep running, so the cleanup the pre-exit hook | ||
| // stood down has to come back. | ||
| await invoke("resume_desktop_update_cleanup").catch(() => {}); |
There was a problem hiding this comment.
On Windows, if resume_desktop_update_cleanup rejects because restoring the job limit fails, this handler suppresses that result and the outer recovery path starts the backend anyway. The app then continues with kill-on-close disabled, so a later force-exit can orphan the restarted backend and an update retry can again reach the installer with live environment users. The current revision's fresh evidence is that the new native resume command explicitly returns these job-object errors, but both frontend recovery call sites discard them; keep the backend stopped or surface a fatal recovery state when re-arming fails.
Useful? React with 👍 / 👎.
terminate_all now applies the same test the startup sweep does: a pid whose identity cannot be read is left alone and kept in the record for the next launch to retry, rather than signalled on the chance it is still ours. The startup reaper re-checks liveness before dropping a record, so a kill that did not take stays reapable, and the breadcrumb unlink happens under the record lock so a concurrent adopt cannot have its record deleted from under it.
6c1efba to
ba665d4
Compare
Linux startup probes prctl with the read-only PR_GET_PDEATHSIG before reporting the parent-death signal as in force, so a seccomp or container policy that blocks it is reported as such rather than as a guarantee nothing keeps. The backstop sweep takes its snapshot under the lock the writes already hold.
ba665d4 to
25910c2
Compare
for more information, see https://pre-commit.ci
…ts leader The DiffusionGemma runner and sd-cli were spawned with lifetime kwargs that are empty on macOS and never recorded, so nothing could reap them. Both now adopt at spawn. The Windows tool capture also revalidated a pid that no longer exists once the wrapper exits, which is the case it was added for; each tool tree gets its own job object instead, with the pid path as the fallback. The startup breadcrumb sweep takes the tree too, not just the leader. Probe PR_SET_PDEATHSIG itself, since seccomp can filter prctl per operation, and stop the post-update backend restart when kill-on-close cannot be re-armed.
for more information, see https://pre-commit.ci
25910c2 to
ba27737
Compare
The runner shares Studio's process group, so the startup sweep could only signal the runner itself and its visual server kept the GPU. start_new_session makes it a group leader, which is what _posix_terminate needs to killpg. Skip and Restart is offered on every error, so gating only the recovery path still let a user start a backend while kill-on-close was disarmed. The flag moved to a ref both paths check.
for more information, see https://pre-commit.ci
ba27737 to
952c7e0
Compare
A recorded leader can exit first and leave its group running, and the sweep then skipped the entry and deleted the record. The child's own process group is recorded at adopt time (only when it leads one, never Studio's) and signalled when the leader is gone. The group id is the dead leader's pid, which the kernel holds while any task still references it as a group, so it cannot belong to anyone else. comm is mutable, so a worker calling prctl(PR_SET_NAME) or setproctitle read as a recycled pid and was dropped unsignalled. Identity is the start time alone; records written with starttime:comm still compare equal.
952c7e0 to
35026b0
Compare
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 35026b0d73
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| timeout = 15, | ||
| creationflags = getattr(subprocess, "CREATE_NO_WINDOW", 0), | ||
| ) | ||
| return |
There was a problem hiding this comment.
On Windows, if the recorded leader exits after _reap_one_record verifies it but before taskkill runs, taskkill reports failure while subprocess.run still returns normally because check defaults to false (Python documentation). This unconditional return skips the fallback; the later liveness check sees the leader gone, leaves unresolved false, and deletes the only record of descendants that /T could no longer reach, permanently orphaning them. Inspect returncode and preserve the record or run the fallback when the command fails.
Useful? React with 👍 / 👎.
forget_pid dropped the recorded group as soon as its leader was reaped, so a shim exiting before its visual server took the only handle on that group with it. The record is kept while the group still has members, and the backstop then reaps it. The Windows backstop now takes the tree, matching the startup sweep, and the identity probe prototypes CloseHandle like every other handle-width call in this module.
35026b0 to
148876a
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f870e13bd8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if !crate::windows_job::kill_on_close_armed().unwrap_or(false) { | ||
| crate::windows_job::resume_after_update_installer().map_err(|error| { | ||
| format!("Refusing to start the backend with crash cleanup disarmed: {error}") | ||
| })?; |
There was a problem hiding this comment.
Fresh evidence in this revision is that the native start guard re-enables only the Windows job limit, not the termination guard consumed by the updater hook. If a webview reload or tray start reaches this path after the job was suspended, the backend starts successfully, but TERMINATION_CLEANUP remains true; on the next update, cleanup_child_processes therefore skips the live backend before suspend_for_update_installer disables kill-on-close, allowing the installer to run while the backend still holds the environment open. Reset the termination guard together with the job limit here.
Useful? React with 👍 / 👎.
| # spawned it. Only those deferred records are revisited, so nothing is | ||
| # signalled twice. | ||
| pending = sorted(directory.glob("*.json")) | ||
| for _pass in range(4): |
There was a problem hiding this comment.
Fresh evidence in this revision is the new retry loop, but its hard four-pass limit still leaves descendants behind when there are five or more dependent owner records ordered child-before-owner. Each pass can kill only the next owner in that chain, so the remaining live owner's record is still deferred when the loop ends and its sidecars survive until another Studio launch. Continue while a pass kills an owner and deferred records remain rather than imposing an arbitrary depth limit.
Useful? React with 👍 / 👎.
for more information, see https://pre-commit.ci
…med record Re-enabling the Windows job on the spawn path left TERMINATION_CLEANUP set, so the next update attempt read cleanup as armed, skipped the resume, and its pre-exit hook suspended kill-on-close without stopping the backend first. An identity read back from a record can be any JSON value; reaching split() with a number raised through the whole startup sweep, so one bad file left every other orphan running. Non-strings now read as unverifiable, which keeps the pid unsignalled rather than trusted. A group whose members exited on the SIGTERM keeps answering killpg(pgid, 0) where pid 1 does not reap, so the reap now rechecks membership instead of waiting out the grace period.
f870e13 to
1e7f681
Compare
for more information, see https://pre-commit.ci
The leader-only fallback runs when the job object was unavailable, which is exactly when the record is the only handle on those workers. Both callers read the dead leader as the tree being gone and dropped it, so anything that survived became unreachable. The tree kill now reports whether it took, and a failure keeps the pid tracked and its record on disk for the next launch.
1e7f681 to
b38d961
Compare
for more information, see https://pre-commit.ci
The installer announced its validation server as stopped as soon as the group leader exited, which drops the record while a child that ignored the SIGTERM is still holding the GPU. It now waits for the group to empty, escalates to SIGKILL, and only announces the stop once nothing is left. An installer timeout killed the installer alone and left the announced server for a sweep that never runs while this process lives; those children are now terminated with it. The diffusion group id is kept from the spawn, so a shim that exited before the kill path (a failed health check, a crash before a reload) no longer leaves its visual server with nothing able to reach it.
b38d961 to
d9b125e
Compare
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d9b125e349
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| try: | ||
| if _is_windows(): | ||
| _windows_terminate_tree(pid) | ||
| else: | ||
| _posix_terminate(pid, timeout) | ||
| except Exception: # noqa: BLE001 - best effort, like the rest of this | ||
| pass | ||
| forget_pid(pid) # keeps the record if its group is still up |
There was a problem hiding this comment.
On Windows, when the installer watchdog calls terminate_pid for an announced validation server and taskkill /T /F fails, _windows_terminate_tree returns false specifically because the tree may still be running, but this helper ignores that result and unconditionally forgets the PID. Fresh evidence beyond the earlier reaper-path finding is the new installer-timeout caller in update_flow.py, which uses this helper while the backend remains alive; the surviving validation tree therefore loses its breadcrumb and can keep holding the GPU or staged files through retries. Retain the record unless tree termination is confirmed.
Useful? React with 👍 / 👎.
# Conflicts: # studio/src-tauri/src/main.rs
The timeout path took them; a nonzero exit or a stream that ended mid-line left them running. This process stays up after an update failure, and its own live record shields those pids from a sweep that would not run anyway, so the cleanup now happens in the finally that covers every way out.
for more information, see https://pre-commit.ci
terminate_pid falls back to the recorded process group when the leader has already exited, keeps the record when a Windows tree kill could not be confirmed, drains the announced children under a lock so the watchdog and the reader thread cannot race, and the installer arms the parent-death signal on the validation server it puts in a session of its own.
for more information, see https://pre-commit.ci
d9b125e to
35cef75
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 35cef75347
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if _is_windows(): | ||
| # False is "only the leader was signalled": nothing else names those | ||
| # workers, so the record has to outlive this call. | ||
| tree_stands = not _windows_terminate_tree(pid) | ||
| else: | ||
| _posix_terminate(pid, timeout) |
There was a problem hiding this comment.
When an announced validation server exits before the installer timeout/error cleanup drains it, its PID can be recycled before terminate_pid runs. This path reads only the recorded PGID and invokes _windows_terminate_tree or _posix_terminate without comparing the current process identity to _tracked_pids, so it can terminate an unrelated process (and, on Windows, that process's descendants). Verify the stored identity immediately before signaling, as terminate_all and _reap_one_record already do.
Useful? React with 👍 / 👎.
| if identity is not None or not _pid_alive(pid): | ||
| return identity | ||
| identity = _pid_identity(pid) | ||
| if identity is not None: | ||
| _tracked_pids[pid] = identity |
There was a problem hiding this comment.
If a short-lived child exits before adopt_pid can capture its identity, it remains tracked with identity=None; a later breadcrumb rewrite after that PID has been reused makes _refreshed_identity capture the unrelated process's identity. Subsequent cleanup then regards that identity as verified and signals the unrelated process. Preserve the unverifiable entry as such, or record that the original PID was already dead rather than filling its identity on a later write.
Useful? React with 👍 / 👎.
…dentity before a single-pid kill The desktop stops this backend by signalling its process group and force-kills it five seconds later, so a runner in a session of its own survives a backend that is slow to shut down and keeps the GPU until the next launch sweeps it. Put it back in that group, as the component installer already is, and reach the visual server by walking the runner's children instead of killpg. The cached group id goes with it: a pid is reusable once nothing holds the number as a process group any more, so an id kept past its group eventually names a stranger. terminate_pid signalled on the pid alone. An announced validation server can exit without the line that clears it, so run the same identity test terminate_all does before either termination branch.
35cef75 to
3b294cd
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3b294cda84
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if identity is None or current is None: | ||
| return True |
There was a problem hiding this comment.
When identity collection fails transiently (for example, a macOS ps call times out), this treats an unverifiable PID as the original descendant. If that descendant exits and its PID is reused during the subsequent grace period, the final _still_the_same check also returns true and terminate_descendants sends SIGKILL to the unrelated replacement process. Preserve the fail-closed behavior used by terminate_all and terminate_pid: do not signal when either the recorded or current identity is unavailable.
Useful? React with 👍 / 👎.
The server was started in a session of its own everywhere, but only Linux can pair that with a parent-death signal. On macOS it left the group Studio force-kills while the only record of it is the announcement the backend has yet to read, so a kill in that window orphaned it with nothing able to find it. It stays in the inherited group there, and the kill path only reaches for killpg when the server actually leads a group, so a shared group is never signalled.
3b294cd to
7a2f09a
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7a2f09a747
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| Only where the parent-death signal can be armed with it, so the two always | ||
| agree: the kill path may only killpg a group the server actually leads. | ||
| """ | ||
| return os.name == "posix" and sys.platform.startswith("linux") |
There was a problem hiding this comment.
On Linux under a seccomp or container policy that rejects PR_SET_PDEATHSIG, this still enables start_new_session; _arm_parent_death ignores the failing prctl return value, so the server has neither the death signal nor membership in the backend's process group. If the installer is killed during validation—particularly before its child announcement is consumed—the validation server can keep holding GPU resources and staged files indefinitely. Probe the capability before detaching, or make the pre-exec hook exit when arming the signal fails.
Useful? React with 👍 / 👎.
|
CI replication finished. |
1 participant
Disposable CI run for unslothai/unsloth#8170. Do not merge; closed after CI.