mirror of
https://github.com/unslothai/unsloth.git
synced 2026-07-09 15:58:41 +00:00
Speed up Studio desktop startup (#6742)
* Speed up Studio desktop startup * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address Studio startup review findings * Keep orphaned run cleanup before readiness * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Wasim Yousef Said
GitHub
parent
02540371c2
commit
755da2f155
10 changed files with 636 additions and 141 deletions
|
|
@ -441,9 +441,30 @@ def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None:
|
|||
).start()
|
||||
|
||||
|
||||
def _warm_rag_embedder() -> None:
|
||||
"""Warm RAG embeddings without blocking backend readiness."""
|
||||
try:
|
||||
from storage import rag_db
|
||||
|
||||
if not rag_db.RAG_AVAILABLE:
|
||||
return
|
||||
from core.rag import embeddings
|
||||
|
||||
embeddings.warm()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
|
||||
|
||||
import time as _time
|
||||
|
||||
_lifespan_started = _time.perf_counter()
|
||||
import structlog as _structlog
|
||||
|
||||
_lifespan_log = _structlog.get_logger(__name__)
|
||||
clear_unsloth_compiled_cache()
|
||||
|
||||
# Remove stale .venv_overlay from old versions; switching now uses .venv_t5/.
|
||||
|
|
@ -454,6 +475,11 @@ async def lifespan(app: FastAPI):
|
|||
# Detect hardware first — sets the DEVICE global used everywhere.
|
||||
detect_hardware()
|
||||
|
||||
_lifespan_log.info(
|
||||
"lifespan hardware detection completed in %.1fms",
|
||||
(_time.perf_counter() - _lifespan_started) * 1000,
|
||||
)
|
||||
|
||||
# Apple Silicon with MLX missing => Train/Export are greyed out (chat-only).
|
||||
# Reinstall mlx by name on a background thread (off the critical path) and
|
||||
# re-detect, so a reinstall/update that dropped mlx self-heals. No-op
|
||||
|
|
@ -465,7 +491,13 @@ async def lifespan(app: FastAPI):
|
|||
import structlog as _structlog
|
||||
_structlog.get_logger(__name__).debug("mlx autorepair skipped: %s", _mlx_exc)
|
||||
|
||||
# Reap download workers orphaned by a previous crash before new downloads start.
|
||||
# Reap workers/runs orphaned by a previous crash before new work starts.
|
||||
try:
|
||||
from storage.studio_db import cleanup_orphaned_runs
|
||||
cleanup_orphaned_runs()
|
||||
except Exception as exc:
|
||||
_lifespan_log.warning("cleanup_orphaned_runs failed at startup: %s", exc)
|
||||
|
||||
reap_hub_orphan_workers()
|
||||
|
||||
# llama.cpp probes: capability (MTP support) + freshness (release age).
|
||||
|
|
@ -479,45 +511,23 @@ async def lifespan(app: FastAPI):
|
|||
app.state.llama_cpp_freshness = None
|
||||
_start_llama_cpp_probes_if_enabled(app)
|
||||
|
||||
from storage.studio_db import cleanup_orphaned_runs
|
||||
|
||||
try:
|
||||
cleanup_orphaned_runs()
|
||||
except Exception as exc:
|
||||
import structlog
|
||||
structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc)
|
||||
|
||||
# Same for RAG: fail ingestion jobs stranded mid-ingest by a crash.
|
||||
try:
|
||||
from storage.rag_db import reconcile_orphaned_ingestion_jobs
|
||||
reconcile_orphaned_ingestion_jobs()
|
||||
except Exception as exc:
|
||||
import structlog
|
||||
structlog.get_logger(__name__).warning(
|
||||
"reconcile_orphaned_ingestion_jobs failed at startup: %s", exc
|
||||
)
|
||||
_lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc)
|
||||
|
||||
_start_helper_precache_if_enabled()
|
||||
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
|
||||
|
||||
# Warm the RAG embedder so the first upload skips the cold load. Non-fatal.
|
||||
def _warm_rag_embedder():
|
||||
try:
|
||||
from storage import rag_db
|
||||
|
||||
if not rag_db.RAG_AVAILABLE:
|
||||
return
|
||||
from core.rag import embeddings
|
||||
|
||||
embeddings.warm()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
threading.Thread(target = _warm_rag_embedder, daemon = True).start()
|
||||
|
||||
# Initialize RSA key pair for API key encryption (external providers)
|
||||
# Initialize RSA key pair for API key encryption (external providers).
|
||||
from core.inference.key_exchange import init_key_pair
|
||||
|
||||
init_key_pair()
|
||||
_lifespan_log.info(
|
||||
"lifespan pre-auth setup completed in %.1fms",
|
||||
(_time.perf_counter() - _lifespan_started) * 1000,
|
||||
)
|
||||
|
||||
if storage.ensure_default_admin():
|
||||
bootstrap_pw = storage.get_bootstrap_password()
|
||||
|
|
@ -532,6 +542,11 @@ async def lifespan(app: FastAPI):
|
|||
print("=" * 60 + "\n")
|
||||
else:
|
||||
app.state.bootstrap_password = storage.get_bootstrap_password()
|
||||
|
||||
_lifespan_log.info(
|
||||
"lifespan startup completed in %.1fms",
|
||||
(_time.perf_counter() - _lifespan_started) * 1000,
|
||||
)
|
||||
yield
|
||||
|
||||
from core.inference.llama_http import aclose as _close_llama_http
|
||||
|
|
@ -919,6 +934,21 @@ install_api_error_handlers(app)
|
|||
# ============ Health and System Endpoints ============
|
||||
|
||||
|
||||
@app.get("/api/liveness")
|
||||
async def liveness_check():
|
||||
"""Cheap process liveness for desktop port validation."""
|
||||
return {
|
||||
"status": "alive",
|
||||
"service": "Unsloth UI Backend",
|
||||
"desktop_protocol_version": 1,
|
||||
"desktop_manageability_version": 1,
|
||||
"supports_desktop_auth": True,
|
||||
"supports_desktop_backend_ownership": True,
|
||||
"studio_root_id": _studio_root_id(),
|
||||
**({"desktop_owner": owner} if (owner := _desktop_owner()) else {}),
|
||||
}
|
||||
|
||||
|
||||
@app.get("/api/health")
|
||||
async def health_check(request: Request):
|
||||
"""Liveness plus launcher capability bits; host fingerprint gated on a bearer.
|
||||
|
|
|
|||
|
|
@ -933,6 +933,9 @@ def run_server(
|
|||
"""
|
||||
global _server, _server_thread, _shutdown_event
|
||||
|
||||
boot_started = time.perf_counter()
|
||||
logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port)
|
||||
|
||||
# Reap every child if the parent dies abnormally (terminal close, Task
|
||||
# Manager kill, SIGKILL); must run before any child can spawn.
|
||||
from utils.process_lifetime import initialize_parent_lifetime
|
||||
|
|
@ -984,7 +987,14 @@ def run_server(
|
|||
from threading import Thread, Event
|
||||
import uvicorn
|
||||
|
||||
import_started = time.perf_counter()
|
||||
|
||||
from main import app, setup_frontend, _IS_COLAB
|
||||
|
||||
logger.info(
|
||||
"Imported FastAPI app in %.1fms",
|
||||
(time.perf_counter() - import_started) * 1000,
|
||||
)
|
||||
from utils.paths import ensure_studio_directories
|
||||
|
||||
# Allow local stdio MCP servers on a loopback bind (the user's own machine),
|
||||
|
|
@ -997,6 +1007,11 @@ def run_server(
|
|||
# Create all standard directories on startup.
|
||||
ensure_studio_directories()
|
||||
|
||||
logger.info(
|
||||
"Ensured Studio directories in %.1fms",
|
||||
(time.perf_counter() - boot_started) * 1000,
|
||||
)
|
||||
|
||||
# Auto-find a free port if the requested one is in use.
|
||||
if not _is_port_free(host, port):
|
||||
original_port = port
|
||||
|
|
@ -1060,6 +1075,11 @@ def run_server(
|
|||
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
||||
_install_uvicorn_startup_log_rewrite(host, display_host)
|
||||
|
||||
logger.info(
|
||||
"run_server pre-uvicorn setup completed in %.1fms",
|
||||
(time.perf_counter() - boot_started) * 1000,
|
||||
)
|
||||
|
||||
ready_event = Event()
|
||||
startup_failed = Event()
|
||||
startup_errors = []
|
||||
|
|
@ -1068,6 +1088,10 @@ def run_server(
|
|||
async def startup(self, *args, **kwargs):
|
||||
await super().startup(*args, **kwargs)
|
||||
if getattr(self, "started", False) and not self.should_exit:
|
||||
logger.info(
|
||||
"Uvicorn startup hook completed in %.1fms",
|
||||
(time.perf_counter() - boot_started) * 1000,
|
||||
)
|
||||
ready_event.set()
|
||||
|
||||
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
|
||||
|
|
@ -1150,6 +1174,11 @@ def run_server(
|
|||
_shutdown_event.set()
|
||||
raise
|
||||
|
||||
logger.info(
|
||||
"run_server uvicorn ready after %.1fms",
|
||||
(time.perf_counter() - boot_started) * 1000,
|
||||
)
|
||||
|
||||
_write_pid_file()
|
||||
import atexit
|
||||
|
||||
|
|
|
|||
|
|
@ -354,12 +354,11 @@ function TauriWrapper({ children }: { children: ReactNode }) {
|
|||
);
|
||||
}
|
||||
|
||||
const showApp = status === "running" && desktopAuthReady;
|
||||
const showApp = status === "running";
|
||||
const desktopBooting = status === "running" && !desktopAuthReady;
|
||||
const showInteractiveApp = showApp && desktopAuthReady;
|
||||
const startupStatus = status === "running" ? "starting" : status;
|
||||
const startupProgressDetail =
|
||||
status === "running" && !desktopAuthReady
|
||||
? "Signing in to desktop session..."
|
||||
: progressDetail;
|
||||
const startupProgressDetail = progressDetail;
|
||||
const usesCustomTitlebar = shouldUseCustomWindowTitlebar();
|
||||
const usesNativeMacTitlebar = shouldUseNativeMacWindowTitlebar();
|
||||
const hidesTitlebarSidebar = HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname);
|
||||
|
|
@ -369,12 +368,23 @@ function TauriWrapper({ children }: { children: ReactNode }) {
|
|||
<TauriUpdateLayer isExternalServer={isExternalServer}>
|
||||
<LlamaUpdateBanner
|
||||
positioned={false}
|
||||
enabled={!hidesTitlebarSidebar}
|
||||
enabled={showInteractiveApp && !hidesTitlebarSidebar}
|
||||
/>
|
||||
<DownloadManagerPanel positioned={false} />
|
||||
{showInteractiveApp ? <DownloadManagerPanel positioned={false} /> : null}
|
||||
</TauriUpdateLayer>
|
||||
<NativeIntentDrain />
|
||||
{children}
|
||||
{showInteractiveApp ? <NativeIntentDrain /> : null}
|
||||
{showInteractiveApp ? children : null}
|
||||
{desktopBooting ? (
|
||||
<div className="pointer-events-none fixed inset-x-0 bottom-5 z-[9999] flex justify-center px-4">
|
||||
<div className="absolute inset-x-4 bottom-16 mx-auto flex max-w-[520px] flex-col items-center gap-2 rounded-2xl border border-border/70 bg-background/95 px-6 py-5 text-center shadow-xl">
|
||||
<div className="font-medium text-sm">Preparing Studio</div>
|
||||
<div className="text-muted-foreground text-xs">The local backend is ready. Signing in to your desktop session before loading chats.</div>
|
||||
</div>
|
||||
<div className="rounded-full border border-border/70 bg-background/95 px-4 py-2 text-xs text-muted-foreground shadow-lg">
|
||||
Signing in to desktop session...
|
||||
</div>
|
||||
</div>
|
||||
) : null}
|
||||
</>
|
||||
) : (
|
||||
<StartupScreen
|
||||
|
|
|
|||
|
|
@ -75,8 +75,7 @@ function externalConflictMessage(preflight: DesktopPreflightResult) {
|
|||
: "A Unsloth server for this install is already running from a terminal. Stop that server, or run `unsloth studio update` from that terminal before using the desktop app.";
|
||||
}
|
||||
|
||||
async function waitForManagedServerReady(
|
||||
invoke: TauriInvoke,
|
||||
async function waitForManagedServerPort(
|
||||
getPort: () => number | null,
|
||||
shouldContinue: () => boolean,
|
||||
): Promise<ManagedStartupResult> {
|
||||
|
|
@ -91,15 +90,7 @@ async function waitForManagedServerReady(
|
|||
continue;
|
||||
}
|
||||
|
||||
const healthy = await invoke<boolean>("check_health", { port });
|
||||
if (!shouldContinue()) {
|
||||
return { status: "aborted" };
|
||||
}
|
||||
if (healthy && getPort() === port) {
|
||||
return { status: "ready", port };
|
||||
}
|
||||
|
||||
await wait(MANAGED_STARTUP_POLL_MS);
|
||||
return { status: "ready", port };
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -280,10 +271,9 @@ export function useTauriBackend() {
|
|||
// backend/run.py keeps the 8888-8908 fallback via server-port/TAURI_PORT.
|
||||
await invoke("start_managed_server", { port: 8888 });
|
||||
|
||||
// Wait for the owned backend's server-port event. Don't attach to an
|
||||
// external backend if the managed start doesn't report a port.
|
||||
const startupResult = await waitForManagedServerReady(
|
||||
invoke,
|
||||
// Rust emits server-port only after validating the desktop-owned process.
|
||||
// Treat that as the UI handoff point instead of doing a second health poll.
|
||||
const startupResult = await waitForManagedServerPort(
|
||||
() => portRef.current,
|
||||
() => startingRef.current,
|
||||
);
|
||||
|
|
|
|||
|
|
@ -65,10 +65,18 @@ pub async fn desktop_preflight(
|
|||
shutdown: tauri::State<'_, ShutdownFlag>,
|
||||
diagnostics: tauri::State<'_, DiagnosticsState>,
|
||||
) -> Result<crate::preflight::DesktopPreflightResult, String> {
|
||||
let started = Instant::now();
|
||||
let (result, adopted_watchdog_generation) =
|
||||
crate::preflight::desktop_preflight_result_with_state(state.inner()).await?;
|
||||
diagnostics::record_preflight(&diagnostics, &result);
|
||||
|
||||
info!(
|
||||
"desktop_preflight completed disposition={:?} port={:?} in {}ms",
|
||||
result.disposition,
|
||||
result.port,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
|
||||
if let Some((generation, newly_adopted)) = adopted_watchdog_generation {
|
||||
if newly_adopted {
|
||||
if let Some(port) = result.port {
|
||||
|
|
@ -205,9 +213,17 @@ pub async fn start_managed_server(
|
|||
port: u16,
|
||||
) -> Result<(), String> {
|
||||
info!("start_managed_server command called with port {}", port);
|
||||
|
||||
let started = Instant::now();
|
||||
let diagnostics_state = diagnostics.inner().clone();
|
||||
let generation = process::start_backend(&app, &state, port, &shutdown, &diagnostics_state)?;
|
||||
|
||||
info!(
|
||||
"start_managed_server spawned generation={} in {}ms",
|
||||
generation,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
|
||||
let watchdog_state = state.inner().clone();
|
||||
let watchdog_shutdown = shutdown.inner().clone();
|
||||
let watchdog_app = app.clone();
|
||||
|
|
|
|||
|
|
@ -93,7 +93,7 @@ enum PreviousAppPidStatus {
|
|||
Uncertain,
|
||||
}
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
struct HealthDesktopOwner {
|
||||
kind: Option<String>,
|
||||
token_sha256: Option<String>,
|
||||
|
|
@ -101,15 +101,7 @@ struct HealthDesktopOwner {
|
|||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct HealthResponse {
|
||||
status: Option<String>,
|
||||
service: Option<String>,
|
||||
version: Option<String>,
|
||||
desktop_protocol_version: Option<u16>,
|
||||
desktop_manageability_version: Option<u16>,
|
||||
supports_desktop_auth: Option<bool>,
|
||||
supports_desktop_backend_ownership: Option<bool>,
|
||||
studio_root_id: Option<String>,
|
||||
desktop_owner: Option<HealthDesktopOwner>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
|
|
@ -123,6 +115,18 @@ struct DesktopLoginPayload<'a> {
|
|||
secret: &'a str,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Deserialize)]
|
||||
pub(crate) struct DesktopLiveness {
|
||||
status: Option<String>,
|
||||
service: Option<String>,
|
||||
desktop_protocol_version: Option<u16>,
|
||||
desktop_manageability_version: Option<u16>,
|
||||
supports_desktop_auth: Option<bool>,
|
||||
supports_desktop_backend_ownership: Option<bool>,
|
||||
studio_root_id: Option<String>,
|
||||
desktop_owner: Option<HealthDesktopOwner>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
struct TokenResponse {
|
||||
access_token: String,
|
||||
|
|
@ -290,10 +294,10 @@ impl BackendOwnerState {
|
|||
}
|
||||
|
||||
pub(crate) fn verifies_exact_port_blocking(&self, port: u16) -> bool {
|
||||
match fetch_health_blocking(port) {
|
||||
Ok(Some(health)) => {
|
||||
health_verifies_metadata(&health, &self.metadata)
|
||||
&& lifecycle_control_block_reason(&health).is_none()
|
||||
match fetch_liveness_blocking(port) {
|
||||
Ok(Some(liveness)) => {
|
||||
liveness_verifies_metadata(&liveness, &self.metadata)
|
||||
&& lifecycle_control_block_reason(&liveness).is_none()
|
||||
}
|
||||
_ => false,
|
||||
}
|
||||
|
|
@ -498,66 +502,110 @@ pub(crate) fn test_owner_state(root_id: &str, token: &str, port: u16) -> Backend
|
|||
}
|
||||
}
|
||||
|
||||
fn health_verifies_metadata(health: &HealthResponse, metadata: &DesktopBackendMetadata) -> bool {
|
||||
let healthy = health.status.as_deref() == Some("healthy")
|
||||
&& health.service.as_deref() == Some("Unsloth UI Backend");
|
||||
let Some(owner) = health.desktop_owner.as_ref() else {
|
||||
fn liveness_verifies_metadata(
|
||||
liveness: &DesktopLiveness,
|
||||
metadata: &DesktopBackendMetadata,
|
||||
) -> bool {
|
||||
let alive = matches!(liveness.status.as_deref(), Some("alive") | Some("healthy"))
|
||||
&& liveness.service.as_deref() == Some("Unsloth UI Backend");
|
||||
let Some(owner) = liveness.desktop_owner.as_ref() else {
|
||||
return false;
|
||||
};
|
||||
healthy
|
||||
alive
|
||||
&& owner_matches_metadata(
|
||||
metadata,
|
||||
health.studio_root_id.as_deref(),
|
||||
liveness.studio_root_id.as_deref(),
|
||||
owner.kind.as_deref(),
|
||||
owner.token_sha256.as_deref(),
|
||||
)
|
||||
}
|
||||
|
||||
fn lifecycle_control_block_reason(health: &HealthResponse) -> Option<String> {
|
||||
if health.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) {
|
||||
fn lifecycle_control_block_reason(liveness: &DesktopLiveness) -> Option<String> {
|
||||
if liveness.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) {
|
||||
return Some("desktop_protocol_incompatible".to_string());
|
||||
}
|
||||
if health.supports_desktop_auth != Some(true) {
|
||||
if liveness.supports_desktop_auth != Some(true) {
|
||||
return Some("desktop_auth_unsupported".to_string());
|
||||
}
|
||||
if health.desktop_manageability_version.unwrap_or(0)
|
||||
if liveness.desktop_manageability_version.unwrap_or(0)
|
||||
< crate::preflight::DESKTOP_MANAGEABILITY_VERSION
|
||||
{
|
||||
return Some("desktop_manageability_unsupported".to_string());
|
||||
}
|
||||
if health.supports_desktop_backend_ownership != Some(true) {
|
||||
if liveness.supports_desktop_backend_ownership != Some(true) {
|
||||
return Some("desktop_backend_ownership_unsupported".to_string());
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
fn ready_for_use_status(health: &HealthResponse) -> OwnedBackendReadiness {
|
||||
match crate::preflight::backend_version_stale_reason(health.version.as_deref()) {
|
||||
fn ready_for_use_status(health: Option<&HealthResponse>) -> OwnedBackendReadiness {
|
||||
let version = health
|
||||
.and_then(|h| h.version.as_deref())
|
||||
.filter(|v| !v.is_empty());
|
||||
match crate::preflight::backend_version_stale_reason(version) {
|
||||
Some(reason) => OwnedBackendReadiness::Stale { reason },
|
||||
None => OwnedBackendReadiness::Ready,
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_health(port: u16) -> Result<Option<HealthResponse>, reqwest::Error> {
|
||||
async fn health_ready_status(port: u16) -> OwnedBackendReadiness {
|
||||
match fetch_health(port).await {
|
||||
Ok(health) => ready_for_use_status(health.as_ref()),
|
||||
Err(reason) => OwnedBackendReadiness::Stale { reason },
|
||||
}
|
||||
}
|
||||
|
||||
async fn fetch_liveness(port: u16) -> Result<Option<DesktopLiveness>, reqwest::Error> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(LOCAL_HTTP_TIMEOUT)
|
||||
.build()?;
|
||||
for path in ["/api/liveness", "/api/health"] {
|
||||
let response = client
|
||||
.get(format!("http://127.0.0.1:{port}{path}"))
|
||||
.send()
|
||||
.await?;
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" {
|
||||
continue;
|
||||
}
|
||||
if !response.status().is_success() {
|
||||
return Ok(None);
|
||||
}
|
||||
return response.json::<DesktopLiveness>().await.map(Some);
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
|
||||
fn fetch_liveness_blocking(port: u16) -> Result<Option<DesktopLiveness>, String> {
|
||||
for path in ["/api/liveness", "/api/health"] {
|
||||
let response = http_request_blocking(port, "GET", path, &[], &[])?;
|
||||
if response.status == 404 && path == "/api/liveness" {
|
||||
continue;
|
||||
}
|
||||
if !(200..300).contains(&response.status) {
|
||||
return Ok(None);
|
||||
}
|
||||
return serde_json::from_slice::<DesktopLiveness>(&response.body)
|
||||
.map(Some)
|
||||
.map_err(|e| e.to_string());
|
||||
}
|
||||
Ok(None)
|
||||
}
|
||||
async fn fetch_health(port: u16) -> Result<Option<HealthResponse>, String> {
|
||||
let client = reqwest::Client::builder()
|
||||
.timeout(LOCAL_HTTP_TIMEOUT)
|
||||
.build()
|
||||
.map_err(|e| e.to_string())?;
|
||||
let response = client
|
||||
.get(format!("http://127.0.0.1:{port}/api/health"))
|
||||
.send()
|
||||
.await?;
|
||||
.await
|
||||
.map_err(|e| e.to_string())?;
|
||||
if !response.status().is_success() {
|
||||
return Ok(None);
|
||||
}
|
||||
response.json::<HealthResponse>().await.map(Some)
|
||||
}
|
||||
|
||||
fn fetch_health_blocking(port: u16) -> Result<Option<HealthResponse>, String> {
|
||||
let response = http_request_blocking(port, "GET", "/api/health", &[], &[])?;
|
||||
if !(200..300).contains(&response.status) {
|
||||
return Ok(None);
|
||||
}
|
||||
serde_json::from_slice::<HealthResponse>(&response.body)
|
||||
response
|
||||
.json::<HealthResponse>()
|
||||
.await
|
||||
.map(Some)
|
||||
.map_err(|e| e.to_string())
|
||||
}
|
||||
|
|
@ -618,21 +666,21 @@ pub(crate) async fn probe_owned_backend_state(
|
|||
};
|
||||
let mut verified = Vec::new();
|
||||
for port in ports {
|
||||
let health = match fetch_health(port).await {
|
||||
Ok(Some(health)) => health,
|
||||
let liveness = match fetch_liveness(port).await {
|
||||
Ok(Some(liveness)) => liveness,
|
||||
Ok(None) => continue,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
"Desktop-owned backend probe skipped port {} after health error: {}",
|
||||
"Desktop-owned backend probe skipped port {} after liveness error: {}",
|
||||
port, error
|
||||
);
|
||||
continue;
|
||||
}
|
||||
};
|
||||
if !health_verifies_metadata(&health, &owner.metadata) {
|
||||
if !liveness_verifies_metadata(&liveness, &owner.metadata) {
|
||||
continue;
|
||||
}
|
||||
if let Some(reason) = lifecycle_control_block_reason(&health) {
|
||||
if let Some(reason) = lifecycle_control_block_reason(&liveness) {
|
||||
return OwnedBackendProbe::Unmanageable { port, reason };
|
||||
}
|
||||
if !desktop_login_route_compatible(port).await {
|
||||
|
|
@ -646,7 +694,7 @@ pub(crate) async fn probe_owned_backend_state(
|
|||
return OwnedBackendProbe::Unmanageable { port, reason };
|
||||
}
|
||||
}
|
||||
verified.push((port, ready_for_use_status(&health)));
|
||||
verified.push((port, health_ready_status(port).await));
|
||||
}
|
||||
|
||||
if verified.len() != 1 {
|
||||
|
|
@ -967,12 +1015,11 @@ mod tests {
|
|||
}
|
||||
|
||||
#[test]
|
||||
fn health_verification_requires_root_kind_and_token_sha() {
|
||||
fn liveness_verification_requires_root_kind_and_token_sha() {
|
||||
let metadata = metadata(1, Some(8888));
|
||||
let health = HealthResponse {
|
||||
status: Some("healthy".to_string()),
|
||||
let liveness = DesktopLiveness {
|
||||
status: Some("alive".to_string()),
|
||||
service: Some("Unsloth UI Backend".to_string()),
|
||||
version: Some("2026.5.2".to_string()),
|
||||
desktop_protocol_version: Some(1),
|
||||
desktop_manageability_version: Some(1),
|
||||
supports_desktop_auth: Some(true),
|
||||
|
|
@ -983,12 +1030,12 @@ mod tests {
|
|||
token_sha256: Some(token_sha256(TOKEN)),
|
||||
}),
|
||||
};
|
||||
assert!(health_verifies_metadata(&health, &metadata));
|
||||
assert!(liveness_verifies_metadata(&liveness, &metadata));
|
||||
|
||||
let mut wrong_root = health;
|
||||
let mut wrong_root = liveness;
|
||||
wrong_root.studio_root_id =
|
||||
Some("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string());
|
||||
assert!(!health_verifies_metadata(&wrong_root, &metadata));
|
||||
assert!(!liveness_verifies_metadata(&wrong_root, &metadata));
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
|
|
|
|||
|
|
@ -193,6 +193,9 @@ pub async fn desktop_preflight_result_with_state(
|
|||
|
||||
if let Some(snapshot) = crate::process::owned_backend_snapshot(state)? {
|
||||
let Some(owner) = snapshot.owner.clone() else {
|
||||
// TAURI_PORT is emitted only after uvicorn lifespan completes; keep
|
||||
// this ownerless path on full health so auth/bootstrap are ready.
|
||||
|
||||
let probe = match snapshot.port {
|
||||
Some(port) => backend::probe_ownerless_spawned_backend(port).await,
|
||||
None => backend,
|
||||
|
|
@ -494,9 +497,22 @@ mod tests {
|
|||
FakeCli { bin, dir }
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
fn remove_managed_capability_cache() {
|
||||
let _ = std::fs::remove_file(
|
||||
dirs::home_dir()
|
||||
.unwrap()
|
||||
.join(".unsloth")
|
||||
.join("studio")
|
||||
.join("desktop_capability_cache.json"),
|
||||
);
|
||||
}
|
||||
|
||||
#[cfg(unix)]
|
||||
#[tokio::test]
|
||||
async fn managed_cli_capability_probe_classifies_core_cases() {
|
||||
remove_managed_capability_cache();
|
||||
|
||||
for (name, script, stale_reason) in [
|
||||
(
|
||||
"cap-missing",
|
||||
|
|
|
|||
|
|
@ -3,7 +3,10 @@ use super::version::{
|
|||
backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
use log::info;
|
||||
use std::time::Duration;
|
||||
use std::time::Instant;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct DesktopOwnerHealth {
|
||||
|
|
@ -24,6 +27,7 @@ pub(super) struct BackendHealth {
|
|||
}
|
||||
|
||||
pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Option<BackendHealth> {
|
||||
let started = Instant::now();
|
||||
let url = format!("http://127.0.0.1:{port}/api/health");
|
||||
let response = client.get(url).send().await.ok()?;
|
||||
if !response.status().is_success() {
|
||||
|
|
@ -40,6 +44,14 @@ pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Optio
|
|||
.and_then(|v| v.as_str())
|
||||
.map(|s| s == "Unsloth UI Backend")
|
||||
.unwrap_or(false);
|
||||
info!(
|
||||
"Desktop preflight: health probe on port {} healthy={} service={} in {}ms",
|
||||
port,
|
||||
healthy,
|
||||
service,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
|
||||
if !healthy || !service {
|
||||
return None;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,14 +2,30 @@ use super::types::ManagedProbe;
|
|||
use super::version::{
|
||||
backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION,
|
||||
};
|
||||
use serde::Deserialize;
|
||||
use log::{info, warn};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Stdio;
|
||||
use std::time::Duration;
|
||||
use std::time::{Duration, Instant, UNIX_EPOCH};
|
||||
use tokio::io::AsyncReadExt;
|
||||
use tokio::process::Command;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
const MANAGED_CAPABILITY_CACHE_SCHEMA: u16 = 2;
|
||||
|
||||
const FNV64_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
|
||||
const FNV64_PRIME: u64 = 0x100000001b3;
|
||||
const HASHED_MARKER_MAX_BYTES: u64 = 64 * 1024;
|
||||
|
||||
const FALLBACK_MARKER_NAMES: &[&str] = &[
|
||||
"pyvenv.cfg",
|
||||
"uv.lock",
|
||||
"requirements.txt",
|
||||
"python.exe",
|
||||
"python",
|
||||
];
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
struct DesktopCapability {
|
||||
desktop_protocol_version: Option<u16>,
|
||||
desktop_manageability_version: Option<u16>,
|
||||
|
|
@ -20,7 +36,225 @@ struct DesktopCapability {
|
|||
version: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize, Serialize)]
|
||||
struct ManagedCapabilityCache {
|
||||
schema: u16,
|
||||
bin_path: String,
|
||||
bin_size: u64,
|
||||
bin_mtime_ms: u64,
|
||||
studio_root_id: Option<String>,
|
||||
marker_path: Option<String>,
|
||||
marker_size: Option<u64>,
|
||||
marker_mtime_ms: Option<u64>,
|
||||
desktop_protocol_version: u16,
|
||||
desktop_manageability_version: u16,
|
||||
capability: DesktopCapability,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct MarkerFingerprint {
|
||||
path: String,
|
||||
size: u64,
|
||||
mtime_ms: u64,
|
||||
content_hash: Option<u64>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq)]
|
||||
struct ManagedBinFingerprint {
|
||||
bin_path: String,
|
||||
bin_size: u64,
|
||||
bin_mtime_ms: u64,
|
||||
studio_root_id: Option<String>,
|
||||
marker_path: Option<String>,
|
||||
marker_size: Option<u64>,
|
||||
marker_mtime_ms: Option<u64>,
|
||||
}
|
||||
|
||||
fn modified_ms(metadata: &fs::Metadata) -> Option<u64> {
|
||||
metadata
|
||||
.modified()
|
||||
.ok()?
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.ok()
|
||||
.and_then(|duration| u64::try_from(duration.as_millis()).ok())
|
||||
}
|
||||
fn hash_bytes(hash: u64, bytes: &[u8]) -> u64 {
|
||||
bytes.iter().fold(hash, |mut next, byte| {
|
||||
next ^= u64::from(*byte);
|
||||
next.wrapping_mul(FNV64_PRIME)
|
||||
})
|
||||
}
|
||||
|
||||
fn marker_content_hash(path: &Path, metadata: &fs::Metadata) -> Option<u64> {
|
||||
if metadata.len() > HASHED_MARKER_MAX_BYTES {
|
||||
return None;
|
||||
}
|
||||
fs::read(path)
|
||||
.ok()
|
||||
.map(|bytes| hash_bytes(FNV64_OFFSET_BASIS, &bytes))
|
||||
}
|
||||
|
||||
fn marker_candidates_for_bin(bin: &Path) -> Vec<PathBuf> {
|
||||
let Some(scripts_dir) = bin.parent() else {
|
||||
return Vec::new();
|
||||
};
|
||||
let Some(venv_dir) = scripts_dir.parent() else {
|
||||
return Vec::new();
|
||||
};
|
||||
let mut out = Vec::new();
|
||||
|
||||
#[cfg(unix)]
|
||||
{
|
||||
if let Ok(lib_dir) = fs::read_dir(venv_dir.join("lib")) {
|
||||
for entry in lib_dir.flatten() {
|
||||
out.push(
|
||||
entry
|
||||
.path()
|
||||
.join("site-packages")
|
||||
.join("unsloth_cli")
|
||||
.join("commands")
|
||||
.join("studio.py"),
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
for marker_name in FALLBACK_MARKER_NAMES {
|
||||
out.push(venv_dir.join(marker_name));
|
||||
out.push(scripts_dir.join(marker_name));
|
||||
}
|
||||
|
||||
out.push(
|
||||
venv_dir
|
||||
.join("Lib")
|
||||
.join("site-packages")
|
||||
.join("unsloth_cli")
|
||||
.join("commands")
|
||||
.join("studio.py"),
|
||||
);
|
||||
out
|
||||
}
|
||||
|
||||
fn managed_bin_fingerprint(bin: &Path) -> Option<ManagedBinFingerprint> {
|
||||
let bin_metadata = fs::metadata(bin).ok()?;
|
||||
let bin_path = bin
|
||||
.canonicalize()
|
||||
.unwrap_or_else(|_| bin.to_path_buf())
|
||||
.to_string_lossy()
|
||||
.into_owned();
|
||||
|
||||
let studio_root_id = crate::desktop_backend_owner::read_expected_studio_root_id();
|
||||
let mut marker_entries: Vec<MarkerFingerprint> = marker_candidates_for_bin(bin)
|
||||
.into_iter()
|
||||
.filter_map(|path| {
|
||||
let metadata = fs::metadata(&path).ok()?;
|
||||
Some(MarkerFingerprint {
|
||||
path: path
|
||||
.canonicalize()
|
||||
.unwrap_or(path.clone())
|
||||
.to_string_lossy()
|
||||
.into_owned(),
|
||||
size: metadata.len(),
|
||||
mtime_ms: modified_ms(&metadata)?,
|
||||
content_hash: marker_content_hash(&path, &metadata),
|
||||
})
|
||||
})
|
||||
.collect();
|
||||
marker_entries.sort_by(|left, right| left.path.cmp(&right.path));
|
||||
let marker_hash = marker_entries
|
||||
.iter()
|
||||
.fold(FNV64_OFFSET_BASIS, |hash, marker| {
|
||||
let next = hash_bytes(hash, marker.path.as_bytes());
|
||||
let next = hash_bytes(next, &marker.size.to_le_bytes());
|
||||
let next = hash_bytes(next, &marker.mtime_ms.to_le_bytes());
|
||||
if let Some(content_hash) = marker.content_hash {
|
||||
hash_bytes(next, &content_hash.to_le_bytes())
|
||||
} else {
|
||||
next
|
||||
}
|
||||
});
|
||||
let marker_path = (!marker_entries.is_empty()).then(|| "markers".to_string());
|
||||
let marker_size = (!marker_entries.is_empty()).then(|| marker_entries.len() as u64);
|
||||
let marker_mtime_ms = (!marker_entries.is_empty()).then_some(marker_hash);
|
||||
|
||||
Some(ManagedBinFingerprint {
|
||||
bin_path,
|
||||
bin_size: bin_metadata.len(),
|
||||
bin_mtime_ms: modified_ms(&bin_metadata)?,
|
||||
studio_root_id,
|
||||
marker_path,
|
||||
marker_size,
|
||||
marker_mtime_ms,
|
||||
})
|
||||
}
|
||||
|
||||
fn capability_cache_path() -> Option<PathBuf> {
|
||||
dirs::home_dir().map(|home| {
|
||||
home.join(".unsloth")
|
||||
.join("studio")
|
||||
.join("desktop_capability_cache.json")
|
||||
})
|
||||
}
|
||||
|
||||
fn cache_matches(cache: &ManagedCapabilityCache, fingerprint: &ManagedBinFingerprint) -> bool {
|
||||
cache.schema == MANAGED_CAPABILITY_CACHE_SCHEMA
|
||||
&& cache.desktop_protocol_version == DESKTOP_PROTOCOL_VERSION
|
||||
&& cache.desktop_manageability_version == DESKTOP_MANAGEABILITY_VERSION
|
||||
&& cache.bin_path == fingerprint.bin_path
|
||||
&& cache.bin_size == fingerprint.bin_size
|
||||
&& cache.bin_mtime_ms == fingerprint.bin_mtime_ms
|
||||
&& cache.studio_root_id == fingerprint.studio_root_id
|
||||
&& cache.marker_path == fingerprint.marker_path
|
||||
&& cache.marker_size == fingerprint.marker_size
|
||||
&& cache.marker_mtime_ms == fingerprint.marker_mtime_ms
|
||||
&& desktop_capability_ready(&cache.capability)
|
||||
}
|
||||
|
||||
fn read_cached_capability(fingerprint: &ManagedBinFingerprint) -> Option<DesktopCapability> {
|
||||
let path = capability_cache_path()?;
|
||||
let bytes = fs::read(path).ok()?;
|
||||
let cache = serde_json::from_slice::<ManagedCapabilityCache>(&bytes).ok()?;
|
||||
if cache_matches(&cache, fingerprint) {
|
||||
Some(cache.capability)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
fn write_cached_capability(fingerprint: &ManagedBinFingerprint, capability: &DesktopCapability) {
|
||||
let Some(path) = capability_cache_path() else {
|
||||
return;
|
||||
};
|
||||
let cache = ManagedCapabilityCache {
|
||||
schema: MANAGED_CAPABILITY_CACHE_SCHEMA,
|
||||
bin_path: fingerprint.bin_path.clone(),
|
||||
bin_size: fingerprint.bin_size,
|
||||
bin_mtime_ms: fingerprint.bin_mtime_ms,
|
||||
studio_root_id: fingerprint.studio_root_id.clone(),
|
||||
marker_path: fingerprint.marker_path.clone(),
|
||||
marker_size: fingerprint.marker_size,
|
||||
marker_mtime_ms: fingerprint.marker_mtime_ms,
|
||||
desktop_protocol_version: DESKTOP_PROTOCOL_VERSION,
|
||||
desktop_manageability_version: DESKTOP_MANAGEABILITY_VERSION,
|
||||
capability: capability.clone(),
|
||||
};
|
||||
if let Some(parent) = path.parent() {
|
||||
if fs::create_dir_all(parent).is_err() {
|
||||
return;
|
||||
}
|
||||
}
|
||||
let Ok(bytes) = serde_json::to_vec_pretty(&cache) else {
|
||||
return;
|
||||
};
|
||||
if let Err(error) = fs::write(&path, bytes) {
|
||||
warn!(
|
||||
"Managed preflight: could not write capability cache: {}",
|
||||
error
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool {
|
||||
let started = Instant::now();
|
||||
let mut cmd = Command::new(bin);
|
||||
cmd.args(args).stdout(Stdio::null()).stderr(Stdio::null());
|
||||
|
||||
|
|
@ -43,20 +277,33 @@ async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool {
|
|||
}
|
||||
|
||||
let Ok(mut child) = cmd.spawn() else {
|
||||
info!(
|
||||
"Managed preflight probe {:?} failed to spawn in {}ms",
|
||||
args,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return false;
|
||||
};
|
||||
|
||||
match tokio::time::timeout(Duration::from_secs(10), child.wait()).await {
|
||||
let ok = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await {
|
||||
Ok(Ok(status)) => status.success(),
|
||||
_ => {
|
||||
let _ = child.kill().await;
|
||||
let _ = child.wait().await;
|
||||
false
|
||||
}
|
||||
}
|
||||
};
|
||||
info!(
|
||||
"Managed preflight probe {:?} finished ok={} in {}ms",
|
||||
args,
|
||||
ok,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
ok
|
||||
}
|
||||
|
||||
async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
|
||||
let started = Instant::now();
|
||||
let mut cmd = Command::new(bin);
|
||||
cmd.args(["studio", "desktop-capabilities", "--json"])
|
||||
.stdout(Stdio::piped())
|
||||
|
|
@ -81,6 +328,10 @@ async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
|
|||
}
|
||||
|
||||
let Ok(mut child) = cmd.spawn() else {
|
||||
info!(
|
||||
"Managed desktop-capabilities probe failed to spawn in {}ms",
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return None;
|
||||
};
|
||||
let Some(mut stdout) = child.stdout.take() else {
|
||||
|
|
@ -92,9 +343,19 @@ async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
|
|||
Err(_) => {
|
||||
let _ = child.kill().await;
|
||||
let _ = child.wait().await;
|
||||
info!(
|
||||
"Managed desktop-capabilities probe timed out in {}ms",
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return None;
|
||||
}
|
||||
_ => {
|
||||
info!(
|
||||
"Managed desktop-capabilities probe exited unsuccessfully in {}ms",
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return None;
|
||||
}
|
||||
_ => return None,
|
||||
}
|
||||
|
||||
let mut output = Vec::new();
|
||||
|
|
@ -102,7 +363,13 @@ async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
|
|||
return None;
|
||||
}
|
||||
|
||||
serde_json::from_slice::<DesktopCapability>(&output).ok()
|
||||
let capability = serde_json::from_slice::<DesktopCapability>(&output).ok();
|
||||
info!(
|
||||
"Managed desktop-capabilities probe finished ok={} in {}ms",
|
||||
capability.is_some(),
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
capability
|
||||
}
|
||||
|
||||
fn desktop_capability_stale_reason(capability: &DesktopCapability) -> Option<String> {
|
||||
|
|
@ -132,18 +399,48 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool {
|
|||
}
|
||||
|
||||
pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe {
|
||||
let started = Instant::now();
|
||||
if !run_cli_probe(&bin, &["-h"]).await {
|
||||
info!(
|
||||
"Managed preflight: cli unusable for {:?} in {}ms",
|
||||
bin,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return ManagedProbe::Stale {
|
||||
bin,
|
||||
reason: "cli_unusable".to_string(),
|
||||
};
|
||||
}
|
||||
|
||||
let capability = probe_cli_capability(&bin).await;
|
||||
if let Some(capability) = capability {
|
||||
if desktop_capability_ready(&capability) {
|
||||
if let Some(fingerprint) = managed_bin_fingerprint(&bin) {
|
||||
if read_cached_capability(&fingerprint).is_some() {
|
||||
info!(
|
||||
"Managed preflight: using cached desktop capability for {:?} in {}ms",
|
||||
bin,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return ManagedProbe::Ready { bin };
|
||||
}
|
||||
}
|
||||
|
||||
let capability = probe_cli_capability(&bin).await;
|
||||
if let Some(capability) = capability {
|
||||
if let Some(fingerprint) = managed_bin_fingerprint(&bin) {
|
||||
write_cached_capability(&fingerprint, &capability);
|
||||
}
|
||||
if desktop_capability_ready(&capability) {
|
||||
info!(
|
||||
"Managed preflight: cli ready for {:?} in {}ms",
|
||||
bin,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return ManagedProbe::Ready { bin };
|
||||
}
|
||||
info!(
|
||||
"Managed preflight: cli stale for {:?} in {}ms",
|
||||
bin,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
return ManagedProbe::Stale {
|
||||
bin,
|
||||
reason: desktop_capability_stale_reason(&capability)
|
||||
|
|
@ -151,6 +448,11 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe {
|
|||
};
|
||||
}
|
||||
|
||||
info!(
|
||||
"Managed preflight: desktop capability probe failed for {:?} in {}ms",
|
||||
bin,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
ManagedProbe::Stale {
|
||||
bin,
|
||||
reason: "desktop_capability_probe_failed".to_string(),
|
||||
|
|
@ -158,10 +460,17 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe {
|
|||
}
|
||||
|
||||
pub(super) async fn probe_managed_install() -> ManagedProbe {
|
||||
match crate::process::find_unsloth_binary() {
|
||||
let started = Instant::now();
|
||||
let result = match crate::process::find_unsloth_binary() {
|
||||
Some(bin) => probe_managed_bin(bin).await,
|
||||
None => ManagedProbe::Missing,
|
||||
}
|
||||
};
|
||||
info!(
|
||||
"Managed preflight: install probe result {:?} in {}ms",
|
||||
result,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
result
|
||||
}
|
||||
|
||||
pub async fn managed_install_ready() -> bool {
|
||||
|
|
|
|||
|
|
@ -735,6 +735,7 @@ pub fn start_backend(
|
|||
}
|
||||
|
||||
async fn generic_backend_health_ok(port: u16) -> bool {
|
||||
let started = std::time::Instant::now();
|
||||
let client = match reqwest::Client::builder()
|
||||
.timeout(Duration::from_secs(2))
|
||||
.build()
|
||||
|
|
@ -745,49 +746,75 @@ async fn generic_backend_health_ok(port: u16) -> bool {
|
|||
return false;
|
||||
}
|
||||
};
|
||||
let response = match client
|
||||
.get(format!("http://127.0.0.1:{port}/api/health"))
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
let mut last_status = None;
|
||||
let mut json = None;
|
||||
for path in ["/api/liveness", "/api/health"] {
|
||||
let response = match client
|
||||
.get(format!("http://127.0.0.1:{port}{path}"))
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(response) => response,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
"Backend port candidate {} failed health request: {}",
|
||||
port, error
|
||||
);
|
||||
return false;
|
||||
}
|
||||
};
|
||||
if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" {
|
||||
last_status = Some(response.status());
|
||||
continue;
|
||||
}
|
||||
if !response.status().is_success() {
|
||||
warn!(
|
||||
"Backend port candidate {} failed health request: {}",
|
||||
port, error
|
||||
"Backend port candidate {} returned HTTP {} from health",
|
||||
port,
|
||||
response.status()
|
||||
);
|
||||
return false;
|
||||
}
|
||||
};
|
||||
if !response.status().is_success() {
|
||||
json = match response.json::<serde_json::Value>().await {
|
||||
Ok(json) => Some(json),
|
||||
Err(error) => {
|
||||
warn!(
|
||||
"Backend port candidate {} returned invalid health JSON: {}",
|
||||
port, error
|
||||
);
|
||||
return false;
|
||||
}
|
||||
};
|
||||
break;
|
||||
}
|
||||
let Some(json) = json else {
|
||||
warn!(
|
||||
"Backend port candidate {} returned HTTP {} from health",
|
||||
port,
|
||||
response.status()
|
||||
last_status
|
||||
.map(|status| status.to_string())
|
||||
.unwrap_or_else(|| "unknown".to_string())
|
||||
);
|
||||
return false;
|
||||
}
|
||||
let json = match response.json::<serde_json::Value>().await {
|
||||
Ok(json) => json,
|
||||
Err(error) => {
|
||||
warn!(
|
||||
"Backend port candidate {} returned invalid health JSON: {}",
|
||||
port, error
|
||||
);
|
||||
return false;
|
||||
}
|
||||
};
|
||||
let healthy = json
|
||||
let live = json
|
||||
.get("status")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s == "healthy")
|
||||
.map(|s| s == "alive" || s == "healthy")
|
||||
.unwrap_or(false);
|
||||
let service = json
|
||||
.get("service")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s == "Unsloth UI Backend")
|
||||
.unwrap_or(false);
|
||||
healthy && service
|
||||
info!(
|
||||
"Backend port candidate {} liveness live={} service={} in {}ms",
|
||||
port,
|
||||
live,
|
||||
service,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
live && service
|
||||
}
|
||||
|
||||
async fn validate_candidate_port(
|
||||
|
|
@ -798,6 +825,7 @@ async fn validate_candidate_port(
|
|||
generation: u64,
|
||||
port: u16,
|
||||
) {
|
||||
let started = std::time::Instant::now();
|
||||
let owner = {
|
||||
let proc = match state.lock() {
|
||||
Ok(proc) => proc,
|
||||
|
|
@ -852,6 +880,14 @@ async fn validate_candidate_port(
|
|||
}
|
||||
};
|
||||
|
||||
info!(
|
||||
"Validated backend port candidate {} valid={} emit={} in {}ms",
|
||||
port,
|
||||
valid,
|
||||
should_emit,
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
|
||||
if should_emit {
|
||||
diagnostics::record_backend_port(&diagnostics_state, &session_id, port);
|
||||
info!("Validated backend port: {}", port);
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue