unsloth/studio/backend/cloudflare_tunnel.py
Daniel Han c6ad59a7d1
Studio: stop leaking child processes on an abnormal exit (#8170)
* Studio: stop leaking child processes on an abnormal exit

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Review fixes: console handler must not touch signals, one child record 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

* Harden the child record against a malformed or older file

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Close the gaps the first pass left

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.

* Give the Popen doubles the pid a real one always has

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fail closed where the answer is not certain

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Never signal a pid the graceful sweep cannot verify

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.

* Only claim the guarantee when it is actually there

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Record every lifetime-bound child, and reach a Windows tree without its 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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Give the diffusion runner its own group, and gate every restart path

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Reap a group whose leader is gone, and drop comm from pid identity

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Keep a group's record until the group is gone

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.

* Put a retained child's group back with it

terminate_all pops the recorded group before deciding what to do, and both
paths that put the pid back dropped it. A leader exiting later then left
nothing able to reach its descendants.

* Believe only confirmed kills in the sweep

taskkill returning nonzero was treated as success, so the documented
single-pid fallback never ran. A group that survived SIGKILL reported
itself resolved, which deleted the last handle on it. And one sweep pass
could skip a worker record whose owner was terminated later in the same
pass, so it repeats while it keeps finding things.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Keep records the sweep could not resolve

A group that survived termination left unresolved false, so the breadcrumb
was deleted with the group still running; the backstop had the same gap. A
zombie owner also read as a live Studio and shielded all of its sidecars.

The retry path re-enters installUpdate and spawns the backend updater, so
every path that starts a child goes through the same re-arm check. The
revisit pass now reconsiders only records deferred for a live owner, so
nothing is signalled twice.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Record tool subprocesses, and do not mistake zombies for a live group

The tool subprocesses lead their own session on POSIX but were never
recorded, so a force quit mid-call left them with nothing able to find
them. They are adopted at spawn and forgotten on confirmed exit.

A leader terminated while alive can leave its group behind, which is the
same loss of the only handle as the dead-leader case. Checked in both the
sweep and the backstop.

killpg(pgid, 0) succeeds for a zombie, so a finished group read as alive
and would have kept its record forever. Membership now ignores zombies.

cloudflared is adopted under the lifecycle lock, before a concurrent stop
can reap it and leave the adoption to record a recycled pid.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Record the component installer, and never wait on a zombie

stream_installer passed the lifetime kwargs but never adopted, and those
kwargs are empty on macOS, so an installer outliving its owner kept
rewriting files under the next launch.

A zombie answers every liveness probe, so terminating one burned the full
grace period per record and reaped nothing. Zombies take the dead-leader
path, which still reaps a group they left behind.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Reset the record lock after a fork, and ask the job whether it is armed

A fork while another thread was inside adopt_pid leaves the child with
the record lock held and no thread to release it, so the next adoption
there blocks forever. The child-side reset already rebuilt the spawner
lock and now rebuilds this one too.

A webview reload rebuilds the update hook with its re-arm flag back at
its initial value while the Windows job can still be disarmed, so the
first gate after a mount reads the job's own limit flags instead.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Retry a child's identity, and treat an unanswerable query as disarmed

A ps that timed out once was recorded as no identity at all, and an
entry without one is never signalled, so that child survived every
later launch. The capture is retried while the process is still there.

On the desktop, a failed desktop_update_cleanup_armed is not evidence
that kill-on-close is in force, so the gate now fails closed and
re-arms.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Start the component installer in its own session, and pin the owner identity

The installer spawns a validation llama-server, and PDEATHSIG reaches
only the direct child while the startup sweep can signal only what the
record names. Leading its own group puts the whole installer tree
within reach of both.

The owner identity is captured once through the same retry a child's
goes through: recorded as None, any process that later reuses the pid
reads as the owner still running and the children in that record are
never reaped. A fork child drops it, since its pid is not the
parent's.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Keep the installer in the group the desktop kills, and heal a missing identity

Reverting the new session from the last round: the desktop stop path
force-kills this backend's process group, and a session of its own took
the installer out of it, so a wedged or SIGKILLed backend left the
installer still rewriting files. Group membership is the stronger
guarantee of the two; the record still names the installer for the
macOS sweep.

A child whose identity could not be read is retried on every breadcrumb
write while it is alive, so a probe that failed once no longer leaves an
entry nothing will ever signal.

A fork child also drops the inherited pid registries: adopting anything
would have written a record claiming its parent's children.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Install the fork reset where the first record is written

It was registered only from the Linux spawn path, which returns before
that everywhere else, so on macOS a fork child kept this process's
children and a record it wrote later claimed them.

* Take the diffusion group down on stop, and do not wait on a zombie

The desktop shutdown only waits for the ordinary stop path, so that is
where the shim's own group has to go: the group is captured before the
wait reaps the leader and killed once the leader is gone. The session
of its own stays, since it is what lets the startup sweep reach a
visual server after a crash.

A group holding nothing but a zombie answers killpg(pgid, 0), so
without a members check every stale record cost the full grace period
where pid 1 does not reap.

A ps that exits nonzero is not an empty group: reporting it as one let
forget_pid drop the only record of a live descendant.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep the re-arm failure message off the Studio branding

* Studio: do not spend the shutdown budget waiting on a child that has exited

An exited child nobody has waited on answers signal 0 exactly like a live
one, so the terminate wait spent its full timeout on a process that was
already gone, once per tracked child and in series. Measured 15.1s for three
of them, 0.16s after this.

The group path keeps waiting while a member is still running, since that is
what holds the GPU. The state read costs a fork off Linux, so it happens
twice a second rather than at the poll rate.

* Studio: answer the group check from the leader instead of scanning every process

Enumerating a process group reads the state of every process on the machine,
which on a busy box is 60ms+, and forget_pid did it on each stop. A leader
that is running already settles the question, and a pid that was never
recorded has neither a group to check nor a record to rewrite. Measured
62ms -> 0.01ms per stop; the leader-has-gone case still falls through to the
scan, which is what finds a child holding the GPU behind it.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: reach the installer's validation server, and never start a backend under a disarmed job

The installer starts a llama-server to validate a build. It is a grandchild
of the backend, so a parent-death signal reaches the installer alone and the
sweep had nothing recording where the server was: an abnormal exit left it
holding the GPU and the staged files. The installer now announces it on
stdout, in its own process group, and the update flow adopts it for as long
as it runs.

On Windows the armed check moves to the path that spawns: the UI gate runs
per update action, but a webview remount can start a backend on its own,
which is exactly the orphan the job object exists to prevent.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: clear the cleanup guard with the re-arm, and survive a malformed 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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep the record when taskkill leaves the tree standing

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: wait for the group, not the leader, before dropping a record

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: make the diffusion group assertion formatting-agnostic

* Studio: stop announced validation servers on any installer exit

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: reach a validation server whose leader or tree kill is gone

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.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Keep the diffusion runner in the group the desktop stops, and check identity 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.

* Keep a macOS validation server in the installer's process group

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.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-08-09 08:06:29 -07:00

976 lines
37 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Free Cloudflare quick tunnel for Unsloth's 0.0.0.0 launches.
The raw http://<ip>:<port> is often unreachable (https-vs-http, blocked ports,
closed security groups); a cloudflared quick tunnel gives a free
https://*.trycloudflare.com URL that works anywhere, with no account or domain.
Best-effort throughout: any failure collapses to "no URL" and Unsloth keeps
running. Stdlib only (back-end imports are lazy) so it is safe to import early.
"""
from __future__ import annotations
import os
import platform
import re
import shutil
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Callable, Optional, Tuple
# cloudflared logs the quick-tunnel URL; match only the URL so we do not depend
# on the surrounding wording, which Cloudflare may change. The negative lookahead
# drops cloudflared's own API host, which appears in failure lines such as
# failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel"
# and must never be mistaken for a usable tunnel URL.
_URL_RE = re.compile(r"https://(?!api\.)[A-Za-z0-9-]+\.trycloudflare\.com")
# cloudflared logs this once per edge connection it establishes. Until at least
# one appears the quick-tunnel URL returns Cloudflare error 1033 (HTTP 530), so
# we wait for it before advertising the URL.
_REGISTERED_MARKER = "Registered tunnel connection"
_RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download"
_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
# A registered edge connection does not mean the hostname resolves yet, so the
# URL is fetched once before it is advertised.
_PUBLIC_PROBE_PATH = "/api/health"
_PUBLIC_PROBE_MARKER = "Unsloth UI Backend"
# One deadline for DNS propagation + the health probe, bounding the startup stall.
_PUBLIC_PROBE_TIMEOUT = 45.0
_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0
_PUBLIC_PROBE_RETRY_DELAY = 1.0
# Wait for the hostname via DoH first: an early OS lookup negative-caches the
# NXDOMAIN for up to 30 min.
_DNS_POLL_DELAY = 2.0
# Retry transient DoH failures, but give up fast when DoH is blocked outright.
_DNS_MAX_DOH_ERRORS = 3
_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A"
# The resolver negative-caches a miss of its own, so a query sent before the
# record can exist blinds the poll for that cache's lifetime. Hold off first.
_DNS_INITIAL_GRACE = 3.0
# A blinded poll cannot recover, so bound its share of the shared deadline.
_DNS_WAIT_MAX = 20.0
# Cloudflare's edge routes by TLS SNI, so it serves the tunnel as soon as the
# connection registers -- before the hostname resolves anywhere. Probing there
# keeps DNS off the startup path entirely.
_EDGE_HOST = "trycloudflare.com"
_EDGE_PROBE_RETRY_DELAY = 0.5
# Bound the wait so the hostname fallback keeps most of the shared deadline.
_EDGE_WAIT_MAX = 15.0
# A network that blocks the edge blocks every attempt, so stop spending the wait.
_EDGE_MAX_UNREACHABLE = 2
def _windows_hidden_kwargs() -> dict:
"""Suppress a child console window on Windows; no-op elsewhere."""
if sys.platform != "win32":
return {}
flags = getattr(subprocess, "CREATE_NO_WINDOW", 0)
return {"creationflags": flags} if flags else {}
def _lifetime_kwargs() -> dict:
"""Bind cloudflared to the parent's lifetime (Linux PDEATHSIG). Lazy +
best-effort so this module still loads standalone (storage_roots-style)."""
try:
from utils.process_lifetime import child_popen_kwargs
return child_popen_kwargs()
except Exception:
return {}
def _adopt_pid(pid: int) -> None:
"""Record cloudflared so a force quit does not strand it (macOS has no
PDEATHSIG). Best-effort, like _lifetime_kwargs above."""
try:
from utils.process_lifetime import adopt_pid
adopt_pid(pid)
except Exception:
pass
def _forget_pid(pid: int) -> None:
try:
from utils.process_lifetime import forget_pid
forget_pid(pid)
except Exception:
pass
def _spawn_child(spawn):
"""Fork on a process-lifetime thread so the PDEATHSIG above means "die with
the parent process", not "die when the worker thread that forked me returns"."""
try:
from utils.process_lifetime import spawn_on_lifetime_thread
except Exception:
return spawn()
return spawn_on_lifetime_thread(spawn)
def _asset_name() -> Optional[Tuple[str, bool]]:
"""(release asset filename, is_tgz) for this OS/arch, or None if unsupported."""
system = platform.system().lower()
machine = platform.machine().lower()
is_x64 = machine in ("x86_64", "amd64", "x64")
is_arm64 = machine in ("aarch64", "arm64")
is_x86 = machine in ("i386", "i686", "x86")
if system == "linux":
if is_x64:
return ("cloudflared-linux-amd64", False)
if is_arm64:
return ("cloudflared-linux-arm64", False)
elif system == "darwin":
if is_arm64:
return ("cloudflared-darwin-arm64.tgz", True)
if is_x64:
return ("cloudflared-darwin-amd64.tgz", True)
elif system == "windows":
if is_x64:
return ("cloudflared-windows-amd64.exe", False)
if is_x86:
return ("cloudflared-windows-386.exe", False)
return None
def _cache_path() -> Optional[Path]:
"""studio_bin_root()/cloudflared(.exe), or None if the studio home is unresolvable."""
try:
from utils.paths.storage_roots import studio_bin_root # lazy: backend-only import
except Exception:
return None
name = "cloudflared.exe" if sys.platform == "win32" else "cloudflared"
return studio_bin_root() / name
def find_cloudflared() -> Optional[str]:
"""Locate an existing cloudflared: PATH first, then the Unsloth bin cache."""
on_path = shutil.which("cloudflared")
if on_path:
return on_path
cached = _cache_path()
if cached is not None and cached.is_file() and os.access(cached, os.X_OK):
return str(cached)
return None
def _download(url: str, dest: Path) -> bool:
"""Download url to dest via urllib (temp file + atomic rename). Best-effort -> bool."""
import tempfile
import urllib.request
tmp_path: Optional[Path] = None
try:
dest.parent.mkdir(parents = True, exist_ok = True)
with tempfile.NamedTemporaryFile(
prefix = dest.name + ".tmp-", dir = dest.parent, delete = False
) as handle:
tmp_path = Path(handle.name)
# GitHub's CDN 403s the default Python-urllib User-Agent.
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _DOWNLOAD_TIMEOUT) as response:
shutil.copyfileobj(response, handle)
if tmp_path.stat().st_size == 0:
raise RuntimeError("empty download")
os.replace(tmp_path, dest)
return True
except Exception:
if tmp_path is not None:
try:
tmp_path.unlink(missing_ok = True)
except Exception:
pass
return False
def _extract_tgz_member(tgz_path: Path, dest: Path) -> bool:
"""Extract just the `cloudflared` member from a darwin .tgz to dest.
Rejects absolute paths and `..` traversal so a hostile archive cannot write
outside dest. Best-effort -> bool.
"""
import tarfile
try:
with tarfile.open(tgz_path, "r:gz") as tar:
member = None
for m in tar.getmembers():
if not m.isfile() or os.path.basename(m.name) != "cloudflared":
continue
if m.name.startswith("/") or ".." in Path(m.name).parts:
continue
member = m
break
if member is None:
return False
src = tar.extractfile(member)
if src is None:
return False
with src, open(dest, "wb") as out:
shutil.copyfileobj(src, out)
return True
except Exception:
return False
def ensure_cloudflared() -> Optional[str]:
"""Return a cloudflared path, downloading + caching the binary once if missing."""
existing = find_cloudflared()
if existing:
return existing
asset = _asset_name()
cached = _cache_path()
if asset is None or cached is None:
return None
name, is_tgz = asset
url = f"{_RELEASE_BASE}/{name}"
try:
cached.parent.mkdir(parents = True, exist_ok = True)
if is_tgz:
tgz = cached.with_suffix(".tgz")
if not _download(url, tgz) or not _extract_tgz_member(tgz, cached):
tgz.unlink(missing_ok = True)
return None
tgz.unlink(missing_ok = True)
elif not _download(url, cached):
return None
if sys.platform != "win32":
os.chmod(cached, 0o755)
return str(cached)
except Exception:
return None
def _wait_for_dns(host: str, deadline: float) -> None:
import json
import urllib.request
now = time.monotonic()
deadline = min(deadline, now + _DNS_WAIT_MAX)
if deadline > now:
time.sleep(min(_DNS_INITIAL_GRACE, deadline - now))
errors = 0
while True:
answered = False
try:
req = urllib.request.Request(
_DOH_URL.format(host = host),
headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"},
)
with urllib.request.urlopen(req, timeout = 5) as response:
answered = bool(json.loads(response.read(65536)).get("Answer"))
errors = 0
except Exception:
errors += 1
if errors >= _DNS_MAX_DOH_ERRORS:
return
if answered:
return
remaining = deadline - time.monotonic()
if remaining <= 0:
return
time.sleep(min(_DNS_POLL_DELAY, remaining))
def _edge_addresses() -> list:
"""Distinct Cloudflare frontends, from a name that resolves before any tunnel exists."""
import socket
addresses = []
try:
resolved = socket.getaddrinfo(_EDGE_HOST, 443, type = socket.SOCK_STREAM)
except Exception:
return addresses
for info in resolved:
address = info[4][0]
# macOS reports the A records as IPv4-mapped under AF_INET6. The mapped
# and bare forms are one frontend, and probing it twice buys nothing.
if address.startswith("::ffff:"):
address = address[len("::ffff:") :]
if address not in addresses:
addresses.append(address)
return addresses[:2]
def _probe_edge(
address: str,
host: str,
timeout: float = _PUBLIC_PROBE_ATTEMPT_TIMEOUT,
) -> Optional[bool]:
"""Ask the edge for the marker as ``host``. None when the edge is unreachable."""
import http.client
import json
import socket
import ssl
request = (
f"GET {_PUBLIC_PROBE_PATH} HTTP/1.1\r\nHost: {host}\r\n"
"User-Agent: unsloth-studio\r\nConnection: close\r\n\r\n"
).encode()
try:
with socket.create_connection((address, 443), timeout = timeout) as raw:
with ssl.create_default_context().wrap_socket(raw, server_hostname = host) as tls:
tls.sendall(request)
response = http.client.HTTPResponse(tls, method = "GET")
response.begin()
body = response.read(4096)
except Exception:
return None
try:
return json.loads(body).get("service") == _PUBLIC_PROBE_MARKER
except Exception:
return False
def _verify_through_edge(host: str, deadline: float) -> bool:
"""Verify at the edge, which selects the tunnel by SNI rather than by address.
Error 1033 and an intercepting proxy's own page are both answers and are not
told apart here, so only the marker ends the wait. Nothing answering at all
is this path being blocked, which the hostname may still get through.
"""
addresses = _edge_addresses()
if not addresses:
return False
deadline = min(deadline, time.monotonic() + _EDGE_WAIT_MAX)
unreachable = 0
while True:
for address in addresses:
remaining = deadline - time.monotonic()
if remaining <= 0:
return False
answer = _probe_edge(address, host, min(_PUBLIC_PROBE_ATTEMPT_TIMEOUT, remaining))
if answer:
return True
unreachable = unreachable + 1 if answer is None else 0
if unreachable >= _EDGE_MAX_UNREACHABLE:
return False
remaining = deadline - time.monotonic()
if remaining <= 0:
return False
time.sleep(min(_EDGE_PROBE_RETRY_DELAY, remaining))
def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool:
import json
import urllib.request
from urllib.parse import urlsplit
deadline = time.monotonic() + timeout
host = urlsplit(url).hostname
if host:
if _verify_through_edge(host, deadline):
return True
# The edge never served the tunnel, so fall back to the hostname and pay
# the DoH wait that keeps an early OS lookup from caching the miss.
_wait_for_dns(host, deadline)
probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}"
while True:
try:
req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response:
body = response.read(4096)
if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER:
return True
except Exception:
pass
remaining = deadline - time.monotonic()
if remaining <= 0:
return False
time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining))
def _process_exited(proc: subprocess.Popen) -> bool:
try:
return proc.poll() is not None
except Exception:
return False
class CloudflareTunnel:
"""A cloudflared quick tunnel to http://localhost:<port>. Best-effort throughout.
Use localhost (not the wildcard bind) as the tunnel origin so cloudflared's
upstream stays local-only.
"""
def __init__(
self,
port: int,
binary: str,
protocol: Optional[str] = None,
):
self.port = port
self.binary = binary
# None lets cloudflared pick its default (quic, with its own http2
# fallback); set to "http2" to force it when quic is blocked.
self.protocol = protocol
self._proc: Optional[subprocess.Popen] = None
self._lock = threading.Lock()
self._stopped = False
self._url_event = threading.Event()
self._ready_event = threading.Event()
self.url: Optional[str] = None
self.ready = False
self.error: Optional[str] = None
self.on_exit: Optional[Callable[["CloudflareTunnel"], None]] = None
self._reader_exited = False
self._runtime_active = False
def start(self) -> None:
cmd = [
self.binary,
"tunnel",
"--url",
f"http://localhost:{self.port}",
"--no-autoupdate",
]
if self.protocol:
cmd += ["--protocol", self.protocol]
with self._lock:
# A stop() that landed before us (e.g. a shutdown in the caller's
# register->start window) marks the tunnel stopped; spawning now would
# orphan a process nobody owns, so refuse.
if self._stopped:
return
_set_studio_tunnel_runtime_active(self, True)
try:
# PDEATHSIG binds to the forking thread, so spawning from the
# settings start worker would kill cloudflared when it returned.
proc = _spawn_child(
lambda: subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
stdin = subprocess.DEVNULL,
text = True,
encoding = "utf-8",
errors = "replace",
bufsize = 1,
**_windows_hidden_kwargs(),
**_lifetime_kwargs(),
)
)
except Exception:
_set_studio_tunnel_runtime_active(self, False)
raise
# Adopted before the lock drops: a stop() that got in first would
# otherwise reap and forget it while nothing was tracked, and this
# would then record whatever inherited the pid.
_adopt_pid(proc.pid)
self._proc = proc
threading.Thread(
target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True
).start()
def _reader(self, proc: subprocess.Popen) -> None:
# Drain cloudflared's output: capture the first trycloudflare URL and the
# first edge-connection registration, and keep draining so it never
# blocks on a full pipe.
try:
if proc.stdout is not None:
for line in proc.stdout:
if self.url is None:
match = _URL_RE.search(line)
if match:
self.url = match.group(0)
self._url_event.set()
if not self.ready and _REGISTERED_MARKER in line:
self.ready = True
self._ready_event.set()
except Exception:
pass
finally:
# stdout closed -> cloudflared has exited. Record why, and unblock any
# waiters at once instead of letting them wait out the full timeout.
if self.url is None:
self.error = "cloudflared exited before emitting a tunnel URL"
elif not self.ready:
self.error = "cloudflared exited before the tunnel connection registered"
else:
self.error = "cloudflared exited"
self._url_event.set()
self._ready_event.set()
with self._lock:
self._reader_exited = True
callback = self.on_exit
if _process_exited(proc):
_set_studio_tunnel_runtime_active(self, False)
if callback is not None:
callback(self)
def wait_for_ready(self, timeout: float = _READY_TIMEOUT) -> Optional[str]:
"""Block until the tunnel is actually serving -- the URL has been minted
*and* at least one edge connection has registered -- or until timeout.
Returns the URL only when ready, so callers never advertise a URL that
would return Cloudflare error 1033 (HTTP 530)."""
self._ready_event.wait(timeout)
return self.url if self.ready else None
def stop(self) -> bool:
"""Terminate the tunnel and report whether process exit was confirmed."""
with self._lock:
# Mark stopped so a start() racing behind us refuses to spawn.
self._stopped = True
proc, self._proc = self._proc, None
if proc is None:
active = _studio_tunnel_runtime_active(self)
if active:
_retain_studio_tunnel_for_stop(self)
return not active
try:
if proc.poll() is None:
proc.terminate()
try:
proc.wait(timeout = 5)
except subprocess.TimeoutExpired:
proc.kill()
try:
proc.wait(timeout = 5)
except Exception:
pass
except Exception:
pass
if _process_exited(proc):
_forget_pid(proc.pid)
_set_studio_tunnel_runtime_active(self, False)
return True
else:
# Preserve both the stop handle and the fail-closed trust state when
# termination could not be confirmed. A later stop can retry.
with self._lock:
if self._proc is None:
self._proc = proc
_retain_studio_tunnel_for_stop(self)
return False
def is_running(self) -> bool:
with self._lock:
proc = self._proc
try:
return proc is not None and proc.poll() is None
except Exception:
return False
def set_on_exit(self, callback: Callable[["CloudflareTunnel"], None]) -> None:
with self._lock:
self.on_exit = callback
reader_exited = self._reader_exited
if reader_exited:
callback(self)
def _publish_if_running(self, callback: Callable[[], None]) -> bool:
with self._lock:
try:
running = (
not self._reader_exited and self._proc is not None and self._proc.poll() is None
)
except Exception:
running = False
if running:
callback()
return running
# Single serving process per Unsloth launch, so one module-level tunnel handle is
# enough; the lock guards the start/stop/shutdown races.
_active_tunnel: Optional[CloudflareTunnel] = None
_active_lock = threading.Lock()
_start_lock = threading.Lock()
# Latched by stop_studio_tunnel so a shutdown landing *between* a start's retry
# attempts aborts the loop instead of starting a tunnel nobody will ever stop.
_shutdown_requested = False
_tunnel_generation = 0
_tunnel_lifecycle = 0
_accepting_starts = True
_tunnel_state = "off"
_tunnel_owner: Optional[str] = None
_tunnel_url: Optional[str] = None
_tunnel_error: Optional[str] = None
_tunnel_port: Optional[int] = None
_tunnel_url_callback: Optional[Callable[[Optional[str]], None]] = None
_tunnel_runtime_callback: Optional[Callable[[bool], None]] = None
_tunnel_runtime_lock = threading.Lock()
_tunnel_runtime_count = 0
_tunnels_pending_stop = set()
_TUNNEL_OWNERS = frozenset({"launch", "settings", "colab"})
def _set_studio_tunnel_runtime_active(tunnel: CloudflareTunnel, active: bool) -> None:
global _tunnel_runtime_count
with _tunnel_runtime_lock:
if not active:
_tunnels_pending_stop.discard(tunnel)
if tunnel._runtime_active == active:
return
tunnel._runtime_active = active
_tunnel_runtime_count += 1 if active else -1
if _tunnel_runtime_callback is not None:
try:
_tunnel_runtime_callback(_tunnel_runtime_count > 0)
except Exception:
pass
def _studio_tunnel_runtime_active(tunnel: CloudflareTunnel) -> bool:
with _tunnel_runtime_lock:
return tunnel._runtime_active
def _retain_studio_tunnel_for_stop(tunnel: CloudflareTunnel) -> None:
with _tunnel_runtime_lock:
if tunnel._runtime_active:
_tunnels_pending_stop.add(tunnel)
def _tunnels_pending_stop_snapshot() -> tuple:
with _tunnel_runtime_lock:
return tuple(_tunnels_pending_stop)
def set_studio_tunnel_runtime_callback(callback: Optional[Callable[[bool], None]]) -> None:
global _tunnel_runtime_callback
with _tunnel_runtime_lock:
_tunnel_runtime_callback = callback
if callback is not None:
try:
callback(_tunnel_runtime_count > 0)
except Exception:
pass
def open_studio_tunnel_lifecycle() -> None:
"""Open a new backend lifecycle and invalidate workers from any prior one."""
global _tunnel_lifecycle, _accepting_starts
with _active_lock:
_tunnel_lifecycle += 1
_accepting_starts = True
def capture_studio_tunnel_start_admission() -> Optional[Tuple[int, int]]:
"""Capture the lifecycle/generation that admitted an asynchronous start."""
with _active_lock:
if not _accepting_starts:
return None
return (_tunnel_lifecycle, _tunnel_generation)
def get_studio_tunnel_control_token() -> Tuple[int, int]:
"""Return the current lifecycle/generation for worker bookkeeping."""
with _active_lock:
return (_tunnel_lifecycle, _tunnel_generation)
def _set_tunnel_url_locked(url: Optional[str]) -> None:
global _tunnel_url
_tunnel_url = url
if _tunnel_url_callback is not None:
try:
_tunnel_url_callback(url)
except Exception:
pass
def set_studio_tunnel_url_callback(callback: Optional[Callable[[Optional[str]], None]]) -> None:
global _tunnel_url_callback
with _active_lock:
_tunnel_url_callback = callback
_set_tunnel_url_locked(_tunnel_url)
def get_studio_tunnel_status() -> dict:
with _active_lock:
return {
"state": _tunnel_state,
"managed_by": _tunnel_owner,
"url": _tunnel_url,
"error": _tunnel_error,
"port": _tunnel_port,
"stop_pending": bool(_tunnels_pending_stop_snapshot()),
}
def _set_failed(generation: int, owner: str, port: int, error: str) -> None:
global _tunnel_state, _tunnel_owner, _tunnel_url, _tunnel_error, _tunnel_port
with _active_lock:
if generation != _tunnel_generation or _shutdown_requested:
return
_tunnel_state = "error"
_tunnel_owner = owner
_set_tunnel_url_locked(None)
_tunnel_error = error
_tunnel_port = port
def _active_tunnel_exited(tunnel: CloudflareTunnel) -> None:
global _active_tunnel, _tunnel_state, _tunnel_owner
global _tunnel_url, _tunnel_error, _tunnel_port
with _active_lock:
if _active_tunnel is not tunnel:
return
if _tunnel_state == "stopping":
return
generation = _tunnel_generation
exit_owner, exit_port = _tunnel_owner, _tunnel_port
exit_error = tunnel.error or "cloudflared exited"
_tunnel_state = "stopping"
_set_tunnel_url_locked(None)
_tunnel_error = None
stopped = tunnel.stop() is not False
with _active_lock:
stopped = stopped or not _studio_tunnel_runtime_active(tunnel)
if not stopped and (_active_tunnel is None or _active_tunnel is tunnel):
_active_tunnel = tunnel
_tunnel_state = "error"
_tunnel_owner = exit_owner
_tunnel_error = "cloudflared could not be stopped"
_tunnel_port = exit_port
elif generation == _tunnel_generation:
_active_tunnel = None
_tunnel_state = "error"
_tunnel_error = exit_error
elif stopped and _active_tunnel is tunnel:
_active_tunnel = None
_tunnel_state = "off"
_tunnel_owner = None
_tunnel_error = None
_tunnel_port = None
def _set_online_locked(url: str) -> None:
global _tunnel_state, _tunnel_url, _tunnel_error
_tunnel_state = "online"
_set_tunnel_url_locked(url)
_tunnel_error = None
def start_studio_tunnel(
port: int,
timeout: float = _READY_TIMEOUT,
*,
managed_by: str = "launch",
admission: Optional[Tuple[int, int]] = None,
) -> Optional[str]:
"""Start a quick tunnel and return its public URL once it is actually
serving, or None (best-effort).
Waits for cloudflared to both mint the URL and register an edge connection,
then fetches /api/health over the public URL, so the caller never advertises
a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host.
If a URL is minted but no connection registers within the window (e.g. quic
is blocked on this network), retries once forcing the http2 protocol. On any
failure the tunnel is stopped and None is returned.
"""
global _active_tunnel, _shutdown_requested, _tunnel_generation
global _tunnel_state, _tunnel_owner, _tunnel_url, _tunnel_error, _tunnel_port
if managed_by not in _TUNNEL_OWNERS:
raise ValueError(f"Unknown Cloudflare tunnel owner: {managed_by}")
with _active_lock:
if not _accepting_starts or _tunnel_state == "stopping" or _tunnels_pending_stop_snapshot():
return None
if admission is not None and admission != (_tunnel_lifecycle, _tunnel_generation):
return None
requested_generation = _tunnel_generation
with _start_lock:
with _active_lock:
if (
_tunnel_state == "online"
and _tunnel_owner == managed_by
and _tunnel_port == port
and _active_tunnel is not None
):
return _tunnel_url
if (
not _accepting_starts
or requested_generation != _tunnel_generation
or _tunnel_state == "stopping"
or _tunnels_pending_stop_snapshot()
or (admission is not None and admission != (_tunnel_lifecycle, _tunnel_generation))
):
return None
_shutdown_requested = False
_tunnel_generation += 1
generation = _tunnel_generation
prior_at_start, _active_tunnel = _active_tunnel, None
_tunnel_state = "starting"
_tunnel_owner = managed_by
_set_tunnel_url_locked(None)
_tunnel_error = None
_tunnel_port = port
if prior_at_start is not None and prior_at_start.stop() is False:
with _active_lock:
if generation == _tunnel_generation:
_active_tunnel = prior_at_start
_tunnel_state = "error"
_tunnel_error = "cloudflared could not be stopped"
return None
binary = ensure_cloudflared()
if not binary:
_set_failed(generation, managed_by, port, "cloudflared is unavailable")
return None
for protocol in (None, "http2"):
with _active_lock:
if _shutdown_requested or generation != _tunnel_generation:
_active_tunnel = None
return None
tunnel = CloudflareTunnel(port, binary, protocol = protocol)
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None and prior.stop() is False:
with _active_lock:
if generation == _tunnel_generation and _active_tunnel is tunnel:
_active_tunnel = prior
_tunnel_state = "error"
_tunnel_error = "cloudflared could not be stopped"
return None
registered = False
try:
tunnel.start()
url = tunnel.wait_for_ready(timeout)
registered = url is not None
if url and not verify_public_url(url):
url = None
except Exception:
url = None
if url:
if hasattr(tunnel, "set_on_exit"):
tunnel.set_on_exit(_active_tunnel_exited)
else:
tunnel.on_exit = _active_tunnel_exited
aborted = False
with _active_lock:
if (
generation != _tunnel_generation
or _shutdown_requested
or _active_tunnel is not tunnel
):
# Stop/shutdown landed while coming up. Detach AND tear
# down: returning this URL would leave a live public
# tunnel no later stop_studio_tunnel() can reach.
aborted = True
was_active = False
if _active_tunnel is tunnel:
_active_tunnel = None
else:
if hasattr(tunnel, "_publish_if_running"):
running = tunnel._publish_if_running(lambda: _set_online_locked(url))
else:
_set_online_locked(url)
running = True
if running:
was_active = True
else:
_active_tunnel = None
_tunnel_state = "error"
_set_tunnel_url_locked(None)
_tunnel_error = tunnel.error or "cloudflared exited"
was_active = False
if aborted or not was_active:
tunnel.stop()
return None
if hasattr(tunnel, "is_running") and not tunnel.is_running():
_active_tunnel_exited(tunnel)
return None
return url
saw_url = tunnel.url is not None
with _active_lock:
was_active = _active_tunnel is tunnel
if was_active:
_tunnel_state = "stopping"
stopped = tunnel.stop() is not False
with _active_lock:
stopped = stopped or not _studio_tunnel_runtime_active(tunnel)
if _shutdown_requested and _active_tunnel is tunnel:
if stopped:
_active_tunnel = None
_tunnel_state = "off"
_tunnel_owner = None
_tunnel_error = None
_tunnel_port = None
else:
_tunnel_state = "error"
_tunnel_error = "cloudflared could not be stopped"
elif generation == _tunnel_generation and _active_tunnel is tunnel:
if stopped:
_active_tunnel = None
# Attempt 1's teardown parked the state at "stopping".
# Leaving it there for the http2 retry makes
# stop_studio_tunnel() early-return and stop nothing.
if _tunnel_state == "stopping" and protocol is None:
_tunnel_state = "starting"
else:
_active_tunnel = tunnel
_tunnel_state = "error"
_tunnel_error = "cloudflared could not be stopped"
if not was_active:
return None
if not stopped:
return None
if not saw_url:
_set_failed(generation, managed_by, port, "cloudflared did not produce a URL")
return None
if registered:
_set_failed(generation, managed_by, port, "Cloudflare URL was not reachable")
return None
_set_failed(generation, managed_by, port, "cloudflared did not register a connection")
return None
def stop_studio_tunnel(*, admission: Optional[Tuple[int, int]] = None) -> None:
"""Terminate the active tunnel, if any. Idempotent."""
global _active_tunnel, _shutdown_requested, _tunnel_generation
global _tunnel_state, _tunnel_owner, _tunnel_url, _tunnel_error, _tunnel_port
with _active_lock:
if admission is not None and admission != (_tunnel_lifecycle, _tunnel_generation):
return
if _tunnel_state == "stopping":
_shutdown_requested = True
_tunnel_generation += 1
return
# Latch so an in-flight start_studio_tunnel won't start a fresh tunnel
# (e.g. its http2 retry) after we have already torn down.
_shutdown_requested = True
_tunnel_generation += 1
stop_generation = _tunnel_generation
tunnel = _active_tunnel
pending = _tunnels_pending_stop_snapshot()
tunnels = list(dict.fromkeys(((tunnel,) if tunnel is not None else ()) + pending))
_tunnel_state = "stopping" if tunnels else "off"
_set_tunnel_url_locked(None)
_tunnel_error = None
for candidate in tunnels:
candidate.stop()
with _active_lock:
if stop_generation == _tunnel_generation or _tunnel_state == "stopping":
pending = _tunnels_pending_stop_snapshot()
if pending:
_active_tunnel = pending[0]
_tunnel_state = "error"
_tunnel_error = "cloudflared could not be stopped"
else:
_active_tunnel = None
_tunnel_state = "off"
_tunnel_owner = None
_tunnel_port = None
def close_studio_tunnel_lifecycle() -> None:
"""Permanently reject queued starts for this backend lifecycle, then stop."""
global _tunnel_lifecycle, _accepting_starts
with _active_lock:
_accepting_starts = False
_tunnel_lifecycle += 1
stop_studio_tunnel()