From 3cd32a282408a1a48cfee6533800852d344f3c11 Mon Sep 17 00:00:00 2001 From: Manan Shah <52329525+Manan17@users.noreply.github.com> Date: Sun, 15 Mar 2026 14:20:48 -0500 Subject: [PATCH] chat only with gguf for mac devices (#4300) * chat only with gguf for mac devices * resolving gpt comments * add change-password for chat only * hide lora adaptors dropdown * solving gpt comments * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * addressing the comment * fixing auth flow --------- Co-authored-by: Datta Nimmaturi Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- cli/commands/studio.py | 18 +++- cli/commands/ui.py | 18 +++- studio/backend/core/inference/defaults.py | 30 ++++++ studio/backend/core/inference/inference.py | 11 +-- studio/backend/core/inference/orchestrator.py | 11 +-- studio/backend/main.py | 5 + studio/backend/routes/models.py | 1 + studio/backend/run.py | 98 ++++++++++++++++--- studio/frontend/src/app/routes/__root.tsx | 10 ++ .../assistant-ui/model-selector.tsx | 38 ++++--- .../assistant-ui/model-selector/pickers.tsx | 8 +- studio/frontend/src/components/navbar.tsx | 16 +-- studio/frontend/src/config/env.ts | 37 +++++++ .../features/auth/components/auth-form.tsx | 37 ++++--- studio/frontend/src/features/auth/session.ts | 5 +- studio/frontend/src/main.tsx | 13 ++- 16 files changed, 283 insertions(+), 73 deletions(-) create mode 100644 studio/backend/core/inference/defaults.py diff --git a/cli/commands/studio.py b/cli/commands/studio.py index 60a0a4d0ff..11a03541a0 100644 --- a/cli/commands/studio.py +++ b/cli/commands/studio.py @@ -106,7 +106,21 @@ def studio_default( args.extend(["--frontend", str(frontend)]) if silent: args.append("--silent") - os.execvp(str(studio_python), args) + # On Windows, os.execvp() spawns a child but the parent lingers, + # so Ctrl+C only kills the parent leaving the child orphaned. + # Use Popen so the parent waits and lets the child's own signal + # handler (in run.py __main__) finish graceful shutdown. + if sys.platform == "win32": + import subprocess as _sp + proc = _sp.Popen(args) + try: + rc = proc.wait() + except KeyboardInterrupt: + # Child's signal handler is doing graceful shutdown — wait + rc = proc.wait() + raise typer.Exit(rc) + else: + os.execvp(str(studio_python), args) else: typer.echo("Studio not set up. Run 'unsloth studio setup' first.") raise typer.Exit(1) @@ -130,6 +144,8 @@ def studio_default( while True: time.sleep(1) except KeyboardInterrupt: + from studio.backend.run import _graceful_shutdown, _server + _graceful_shutdown(_server) typer.echo("\nShutting down...") diff --git a/cli/commands/ui.py b/cli/commands/ui.py index 10db9e243e..1586d507f7 100644 --- a/cli/commands/ui.py +++ b/cli/commands/ui.py @@ -49,7 +49,21 @@ def ui( args.extend(["--frontend", str(frontend)]) if silent: args.append("--silent") - os.execvp(str(studio_python), args) + # On Windows, os.execvp() spawns a child but the parent lingers, + # so Ctrl+C only kills the parent leaving the child orphaned. + # Use Popen so the parent waits and lets the child's own signal + # handler (in run.py __main__) finish graceful shutdown. + if sys.platform == "win32": + import subprocess as _sp + proc = _sp.Popen(args) + try: + rc = proc.wait() + except KeyboardInterrupt: + # Child's signal handler is doing graceful shutdown — wait + rc = proc.wait() + raise typer.Exit(rc) + else: + os.execvp(str(studio_python), args) else: typer.echo("Studio not set up. Run 'unsloth studio setup' first.") raise typer.Exit(1) @@ -73,4 +87,6 @@ def ui( while True: time.sleep(1) except KeyboardInterrupt: + from studio.backend.run import _graceful_shutdown, _server + _graceful_shutdown(_server) typer.echo("\nShutting down...") diff --git a/studio/backend/core/inference/defaults.py b/studio/backend/core/inference/defaults.py new file mode 100644 index 0000000000..c8e23deb09 --- /dev/null +++ b/studio/backend/core/inference/defaults.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Default model lists for inference, split by platform.""" + +import sys + +DEFAULT_MODELS_GGUF = [ + "unsloth/Llama-3.2-1B-Instruct-GGUF", + "unsloth/Llama-3.2-3B-Instruct-GGUF", + "unsloth/Llama-3.1-8B-Instruct-GGUF", + "unsloth/gemma-3-1b-it-GGUF", + "unsloth/gemma-3-4b-it-GGUF", + "unsloth/Qwen3-4B-GGUF", +] + +DEFAULT_MODELS_STANDARD = [ + "unsloth/Qwen3-4B-Instruct-2507", + "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", + "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", + "unsloth/Phi-3.5-mini-instruct", + "unsloth/Gemma-3-4B-it", + "unsloth/Qwen2-VL-2B-Instruct-bnb-4bit", +] + + +def get_default_models() -> list[str]: + if sys.platform == "darwin": + return list(DEFAULT_MODELS_GGUF) + return list(DEFAULT_MODELS_STANDARD) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index cb287bffa5..4b162c5727 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -36,14 +36,9 @@ class InferenceBackend: self.active_model_name = None self.loading_models = set() self.loaded_local_models = [] # [(display_name, path), ...] - self.default_models = [ - "unsloth/Qwen3-4B-Instruct-2507", - "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit", - "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", - "unsloth/Phi-3.5-mini-instruct", - "unsloth/Gemma-3-4B-it", - "unsloth/Qwen2-VL-2B-Instruct-bnb-4bit", - ] + from core.inference.defaults import get_default_models + + self.default_models = get_default_models() self.device = get_device().value self._audio_codec_manager = AudioCodecManager() diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 99c43bbd88..3c87d64b12 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -74,14 +74,9 @@ class InferenceOrchestrator: self.models: dict = {} self.loading_models: set = set() self.loaded_local_models: list = [] - self._static_models = [ - "unsloth/Qwen3-4B-Instruct-2507", - "unsloth/Llama-3.1-8B-Instruct-bnb-4bit", - "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit", - "unsloth/Phi-3.5-mini-instruct", - "unsloth/Gemma-3-4B-it", - "unsloth/Qwen2-VL-2B-Instruct-bnb-4bit", - ] + from core.inference.defaults import get_default_models + + self._static_models = get_default_models() self._top_gguf_cache: Optional[list[str]] = None self._top_gguf_fetched = False diff --git a/studio/backend/main.py b/studio/backend/main.py index 3c2732eb5a..be46272847 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -11,6 +11,7 @@ import os os.environ["PYTHONWARNINGS"] = "ignore" import shutil +import sys import warnings from contextlib import asynccontextmanager @@ -144,10 +145,14 @@ app.include_router(export_router, prefix = "/api/export", tags = ["export"]) @app.get("/api/health") async def health_check(): """Health check endpoint""" + platform_map = {"darwin": "mac", "win32": "windows", "linux": "linux"} + device_type = platform_map.get(sys.platform, sys.platform) + return { "status": "healthy", "timestamp": datetime.now().isoformat(), "service": "Unsloth UI Backend", + "device_type": device_type, } diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 4258731712..586b18718e 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -314,6 +314,7 @@ async def list_models( model_info = ModelDetails( id = model_id, name = model_id.split("/")[-1] if "/" in model_id else model_id, + is_gguf = model_id.upper().endswith("-GGUF"), ) all_models.append(model_info) seen_ids.add(model_id) diff --git a/studio/backend/run.py b/studio/backend/run.py index 4924406d2e..a18d0d95bd 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -92,6 +92,59 @@ def _find_free_port(host: str, start: int, max_attempts: int = 20) -> int: ) +def _graceful_shutdown(server=None): + """Explicitly shut down all subprocess backends and the uvicorn server. + + Called from signal handlers to ensure child processes are cleaned up + before the parent exits. This is critical on Windows where atexit + handlers are unreliable after Ctrl+C. + """ + logger.info("Graceful shutdown initiated — cleaning up subprocesses...") + + # 1. Shut down uvicorn server (releases the listening socket) + if server is not None: + server.should_exit = True + + # 2. Clean up inference subprocess (if instantiated) + try: + from core.inference.orchestrator import _inference_backend + if _inference_backend is not None: + _inference_backend._shutdown_subprocess(timeout = 5.0) + except Exception as e: + logger.warning("Error shutting down inference subprocess: %s", e) + + # 3. Clean up export subprocess (if instantiated) + try: + from core.export.orchestrator import _export_backend + if _export_backend is not None: + _export_backend._shutdown_subprocess(timeout = 5.0) + except Exception as e: + logger.warning("Error shutting down export subprocess: %s", e) + + # 4. Clean up training subprocess (if active) + try: + from core.training.training import _training_backend + if _training_backend is not None: + _training_backend.force_terminate() + except Exception as e: + logger.warning("Error shutting down training subprocess: %s", e) + + # 5. Kill llama-server subprocess (if loaded) + try: + from routes.inference import _llama_cpp_backend + if _llama_cpp_backend is not None: + _llama_cpp_backend._kill_process() + except Exception as e: + logger.warning("Error shutting down llama-server: %s", e) + + logger.info("All subprocesses cleaned up") + + +# The uvicorn server instance — set by run_server(), used by callers +# that need to tell the server to exit (e.g. signal handlers). +_server = None + + def run_server( host: str = "0.0.0.0", port: int = 8000, @@ -106,7 +159,14 @@ def run_server( port: Port to bind to (auto-increments if in use) frontend_path: Path to frontend build directory (optional) silent: Suppress startup messages + + Note: + Signal handlers are NOT registered here so that embedders + (e.g. Colab notebooks) keep their own interrupt semantics. + Standalone callers should register handlers after calling this. """ + global _server + import nest_asyncio nest_asyncio.apply() @@ -138,13 +198,15 @@ def run_server( if not silent: print(f"⚠️ Frontend not found at {frontend_path}") - # Run server + # Create the uvicorn server and expose it for signal handlers + config = uvicorn.Config( + app, host = host, port = port, log_level = "info", access_log = False + ) + _server = uvicorn.Server(config) + + # Run server in a daemon thread def _run(): - config = uvicorn.Config( - app, host = host, port = port, log_level = "info", access_log = False - ) - server = uvicorn.Server(config) - asyncio.run(server.serve()) + asyncio.run(_server.serve()) thread = Thread(target = _run, daemon = True) thread.start() @@ -165,9 +227,11 @@ def run_server( return app -# For direct execution +# For direct execution (also invoked by CLI via os.execvp / subprocess) if __name__ == "__main__": import argparse + import signal + from threading import Event parser = argparse.ArgumentParser(description = "Run Unsloth UI Backend server") parser.add_argument("--host", default = "0.0.0.0", help = "Host to bind to") @@ -187,8 +251,20 @@ if __name__ == "__main__": kwargs["frontend_path"] = Path(args.frontend) run_server(**kwargs) - # Keep running - import time + # ── Signal handler — ensures subprocess cleanup on Ctrl+C ──── + _shutdown_event = Event() + + def _signal_handler(signum, frame): + _graceful_shutdown(_server) + _shutdown_event.set() + + signal.signal(signal.SIGINT, _signal_handler) + signal.signal(signal.SIGTERM, _signal_handler) + + # On Windows, some terminals send SIGBREAK for Ctrl+C / Ctrl+Break + if hasattr(signal, "SIGBREAK"): + signal.signal(signal.SIGBREAK, _signal_handler) + + # Keep running until shutdown signal + _shutdown_event.wait() - while True: - time.sleep(1) diff --git a/studio/frontend/src/app/routes/__root.tsx b/studio/frontend/src/app/routes/__root.tsx index a58343d0c5..0725d383c7 100644 --- a/studio/frontend/src/app/routes/__root.tsx +++ b/studio/frontend/src/app/routes/__root.tsx @@ -2,16 +2,26 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { Navbar } from "@/components/navbar"; +import { usePlatformStore } from "@/config/env"; import { Outlet, createRootRoute, + redirect, useRouterState, } from "@tanstack/react-router"; import { AnimatePresence, motion } from "motion/react"; import { Suspense } from "react"; import { AppProvider } from "../provider"; +const CHAT_ONLY_ALLOWED = new Set(["/", "/chat", "/login", "/signup", "/change-password"]); + export const Route = createRootRoute({ + beforeLoad: ({ location }) => { + const chatOnly = usePlatformStore.getState().isChatOnly(); + if (chatOnly && !CHAT_ONLY_ALLOWED.has(location.pathname)) { + throw redirect({ to: "/chat" }); + } + }, component: RootLayout, }); diff --git a/studio/frontend/src/components/assistant-ui/model-selector.tsx b/studio/frontend/src/components/assistant-ui/model-selector.tsx index 28f94dbd19..2df72139d8 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector.tsx @@ -9,6 +9,7 @@ import { PopoverTrigger, } from "@/components/ui/popover"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { usePlatformStore } from "@/config/env"; import { cn } from "@/lib/utils"; import { ArrowDown01Icon, @@ -111,6 +112,7 @@ function ModelSelectorContent({ dataTour?: string; }) { const hasSelection = Boolean(value); + const chatOnly = usePlatformStore((s) => s.isChatOnly()); return ( - - - Hub models - Fine-tuned - + {chatOnly ? ( + + ) : ( + + + Hub models + Fine-tuned + - - - + + + - - - - + + + + + )} {hasSelection && onEject ? (
diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index 76a32b6129..fe4226c4af 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -11,6 +11,7 @@ import { import { listCachedGguf, listGgufVariants } from "@/features/chat/api/chat-api"; import type { CachedGgufRepo } from "@/features/chat/api/chat-api"; import type { GgufVariantDetail } from "@/features/chat/types/api"; +import { usePlatformStore } from "@/config/env"; import { useDebouncedValue, useGpuInfo, @@ -377,12 +378,15 @@ export function HubModelPicker({ const showHfSection = debouncedQuery.trim().length > 0; const recommendedSet = useMemo(() => new Set(recommendedIds), [recommendedIds]); + const chatOnly = usePlatformStore((s) => s.isChatOnly()); + const hfIds = useMemo(() => { if (!showHfSection) return []; return results .map((result) => result.id) - .filter((id) => !recommendedSet.has(id)); - }, [recommendedSet, results, showHfSection]); + .filter((id) => !recommendedSet.has(id)) + .filter((id) => !chatOnly || isGgufRepo(id)); + }, [recommendedSet, results, showHfSection, chatOnly]); const metricsById = useMemo( () => diff --git a/studio/frontend/src/components/navbar.tsx b/studio/frontend/src/components/navbar.tsx index 2fd2e12088..e778b013f2 100644 --- a/studio/frontend/src/components/navbar.tsx +++ b/studio/frontend/src/components/navbar.tsx @@ -26,6 +26,7 @@ import { } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useTrainingRuntimeStore } from "@/features/training"; +import { usePlatformStore } from "@/config/env"; import { Link, useRouterState } from "@tanstack/react-router"; import { motion } from "motion/react"; import { useState } from "react"; @@ -50,6 +51,8 @@ export function Navbar() { const isTrainingRunning = useTrainingRuntimeStore((s) => s.isTrainingRunning); const [mobileOpen, setMobileOpen] = useState(false); + const chatOnly = usePlatformStore((s) => s.isChatOnly()); + const tourId = getTourId(pathname); const openTour = () => { @@ -63,10 +66,7 @@ export function Navbar() {
{/* Left: logo */} - + Unsloth boolean; +} + +export const usePlatformStore = create()((_, get) => ({ + deviceType: "linux", + fetched: false, + isChatOnly: () => get().deviceType === "mac", +})); + +export async function fetchDeviceType(): Promise { + const { fetched } = usePlatformStore.getState(); + if (fetched) return usePlatformStore.getState().deviceType; + + try { + const res = await fetch("/api/health"); + if (res.ok) { + const data = (await res.json()) as { device_type?: string }; + const deviceType = data.device_type ?? "linux"; + usePlatformStore.setState({ deviceType, fetched: true }); + return deviceType; + } + } catch (err) { + console.warn("[platform] Failed to fetch device type, will retry", err); + } + + return usePlatformStore.getState().deviceType; +} diff --git a/studio/frontend/src/features/auth/components/auth-form.tsx b/studio/frontend/src/features/auth/components/auth-form.tsx index d51fc2814c..4433f29bd5 100644 --- a/studio/frontend/src/features/auth/components/auth-form.tsx +++ b/studio/frontend/src/features/auth/components/auth-form.tsx @@ -89,20 +89,9 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null { let canceled = false; async function initializeAuthForm(): Promise { - if (hasRefreshToken()) { - const refreshed = await refreshSession(); - if (refreshed) { - if (!canceled) setStatusLoading(false); - navigate({ to: getPostAuthRoute() }); - return; - } - } - if (hasAuthToken()) { - if (!canceled) setStatusLoading(false); - navigate({ to: getPostAuthRoute() }); - return; - } - + // Always check the server first — localStorage flags can be stale + // (e.g. tokens from a previous install attempt). The server's + // /api/auth/status is the source of truth for requires_password_change. try { const response = await fetch("/api/auth/status"); if (!response.ok) throw new Error("Failed to load auth status."); @@ -111,6 +100,8 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null { setInitialized(result.initialized); setUsername(result.default_username); setRequiresPasswordChange(result.requires_password_change); + + // Redirect between login ↔ change-password based on server state if (mode === "login" && result.requires_password_change) { navigate({ to: "/change-password" }); return; @@ -119,6 +110,24 @@ export function AuthForm({ mode }: AuthFormProps): ReactElement | null { navigate({ to: "/login" }); return; } + + // On login page, if user already has a valid session and no + // password change is required, skip straight to the app. + if (isLoginMode && !result.requires_password_change) { + if (hasRefreshToken()) { + const refreshed = await refreshSession(); + if (refreshed) { + if (!canceled) setStatusLoading(false); + navigate({ to: getPostAuthRoute() }); + return; + } + } + if (hasAuthToken()) { + if (!canceled) setStatusLoading(false); + navigate({ to: getPostAuthRoute() }); + return; + } + } } } catch (err: unknown) { if (!canceled) { diff --git a/studio/frontend/src/features/auth/session.ts b/studio/frontend/src/features/auth/session.ts index 5ee2a6b13b..3d3502e073 100644 --- a/studio/frontend/src/features/auth/session.ts +++ b/studio/frontend/src/features/auth/session.ts @@ -1,12 +1,14 @@ // SPDX-License-Identifier: AGPL-3.0-only // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +import { usePlatformStore } from "@/config/env"; + export const AUTH_TOKEN_KEY = "unsloth_auth_token"; export const AUTH_REFRESH_TOKEN_KEY = "unsloth_auth_refresh_token"; export const ONBOARDING_DONE_KEY = "unsloth_onboarding_done"; export const AUTH_MUST_CHANGE_PASSWORD_KEY = "unsloth_auth_must_change_password"; -type PostAuthRoute = "/onboarding" | "/studio" | "/change-password"; +type PostAuthRoute = "/onboarding" | "/studio" | "/change-password" | "/chat"; function canUseStorage(): boolean { return typeof window !== "undefined"; @@ -77,5 +79,6 @@ export function resetOnboardingDone(): void { export function getPostAuthRoute(): PostAuthRoute { if (mustChangePassword()) return "/change-password"; + if (usePlatformStore.getState().isChatOnly()) return "/chat"; return isOnboardingDone() ? "/studio" : "/onboarding"; } diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx index 77186e09a2..0db37fdb6e 100644 --- a/studio/frontend/src/main.tsx +++ b/studio/frontend/src/main.tsx @@ -5,6 +5,7 @@ import { StrictMode } from "react"; import { createRoot } from "react-dom/client"; import "./index.css"; +import { fetchDeviceType } from "./config/env"; import { App } from "./app/app"; const globalCrypto = globalThis.crypto as Crypto | undefined; @@ -32,8 +33,10 @@ if (!rootElement) { throw new Error("Root element not found"); } -createRoot(rootElement).render( - - - , -); +fetchDeviceType().then(() => { + createRoot(rootElement).render( + + + , + ); +});