From 37075c542258e87bb556f6ed7496ea88a818c3b6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 07:49:59 -0700 Subject: [PATCH 001/402] Bump install.sh / install.ps1 pin to unsloth>=2026.7.1 (#6943) Co-authored-by: danielhanchen --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 8c667df079..9114f80af9 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2155,7 +2155,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2169,7 +2169,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2235,7 +2235,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2247,7 +2247,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2275,7 +2275,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index 14fbba478d..796d80e401 100755 --- a/install.sh +++ b/install.sh @@ -2706,7 +2706,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2721,7 +2721,7 @@ if [ "$_MIGRATED" = true ]; then # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2925,7 +2925,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2943,7 +2943,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" + --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2975,7 +2975,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." From 07ecdb34c092cb0107f74dbf117616e320ec6ff2 Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Tue, 7 Jul 2026 22:24:32 +0530 Subject: [PATCH 002/402] Sort chat recents by last activity (#6844) * show chat by by last activity * Update chat thread updated_at logic and enhance sidebar chat item handling * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/routes/chat_history.py | 4 +- studio/backend/storage/studio_db.py | 77 +++++++++- .../tests/test_chat_history_storage.py | 134 ++++++++++++++++++ .../frontend/src/components/app-sidebar.tsx | 16 ++- .../chat/hooks/use-chat-sidebar-items.ts | 26 +++- studio/frontend/src/features/chat/types.ts | 1 + .../chat/utils/chat-history-storage.ts | 5 +- studio/frontend/src/i18n/locales/en.ts | 1 + studio/frontend/src/i18n/locales/zh-CN.ts | 1 + 9 files changed, 252 insertions(+), 13 deletions(-) diff --git a/studio/backend/routes/chat_history.py b/studio/backend/routes/chat_history.py index 963d584303..7a27a58a52 100644 --- a/studio/backend/routes/chat_history.py +++ b/studio/backend/routes/chat_history.py @@ -56,6 +56,7 @@ class ChatThread(BaseModel): projectId: Optional[str] = None archived: bool = False createdAt: int + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None forkedFromThreadId: Optional[str] = None @@ -70,6 +71,7 @@ class ChatThreadPatch(BaseModel): projectId: Optional[str] = None archived: Optional[bool] = None createdAt: Optional[int] = None + updatedAt: Optional[int] = None openaiCodeExecContainerId: Optional[str] = None anthropicCodeExecContainerId: Optional[str] = None @@ -252,7 +254,7 @@ async def patch_thread( current_subject: str = Depends(get_current_subject), ): patch = payload.model_dump(exclude_unset = True) - for field in ("title", "modelType", "modelId", "archived", "createdAt"): + for field in ("title", "modelType", "modelId", "archived", "createdAt", "updatedAt"): if field in patch and patch[field] is None: raise HTTPException(status_code = 400, detail = f"{field} cannot be null") if patch.get("projectId") and get_chat_project(patch["projectId"]) is None: diff --git a/studio/backend/storage/studio_db.py b/studio/backend/storage/studio_db.py index 41a9adcc29..87aa50ee26 100644 --- a/studio/backend/storage/studio_db.py +++ b/studio/backend/storage/studio_db.py @@ -240,6 +240,7 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: project_id TEXT, archived INTEGER NOT NULL DEFAULT 0, created_at INTEGER NOT NULL, + updated_at INTEGER, openai_code_exec_container_id TEXT, anthropic_code_exec_container_id TEXT, forked_from_thread_id TEXT, @@ -261,6 +262,24 @@ def _ensure_schema(conn: sqlite3.Connection) -> None: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_thread_id TEXT") if "forked_from_message_id" not in chat_thread_cols: conn.execute("ALTER TABLE chat_threads ADD COLUMN forked_from_message_id TEXT") + if "updated_at" not in chat_thread_cols: + conn.execute("ALTER TABLE chat_threads ADD COLUMN updated_at INTEGER") + # Floor at created_at: forked threads copy older ancestor messages, + # so the fork's creation time must win over the branch message times. + conn.execute( + """ + UPDATE chat_threads SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + """ + ) conn.execute( """ CREATE TABLE IF NOT EXISTS chat_messages ( @@ -992,6 +1011,9 @@ def _chat_thread_from_row(row: sqlite3.Row) -> dict: "projectId": data.get("project_id") or None, "archived": bool(data["archived"]), "createdAt": data["created_at"], + "updatedAt": data.get("updated_at") + if data.get("updated_at") is not None + else data["created_at"], "openaiCodeExecContainerId": data.get("openai_code_exec_container_id"), "anthropicCodeExecContainerId": data.get("anthropic_code_exec_container_id"), "forkedFromThreadId": data.get("forked_from_thread_id"), @@ -1039,8 +1061,8 @@ def upsert_chat_thread(thread: dict) -> dict: conn.execute( """ INSERT INTO chat_threads - (id, title, model_type, model_id, pair_id, project_id, archived, created_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + (id, title, model_type, model_id, pair_id, project_id, archived, created_at, updated_at, openai_code_exec_container_id, anthropic_code_exec_container_id, forked_from_thread_id, forked_from_message_id) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) ON CONFLICT(id) DO UPDATE SET title = excluded.title, model_type = excluded.model_type, @@ -1049,6 +1071,7 @@ def upsert_chat_thread(thread: dict) -> dict: project_id = excluded.project_id, archived = excluded.archived, created_at = excluded.created_at, + updated_at = COALESCE(excluded.updated_at, chat_threads.updated_at), openai_code_exec_container_id = excluded.openai_code_exec_container_id, anthropic_code_exec_container_id = excluded.anthropic_code_exec_container_id, forked_from_thread_id = excluded.forked_from_thread_id, @@ -1063,6 +1086,7 @@ def upsert_chat_thread(thread: dict) -> dict: thread.get("projectId"), 1 if thread.get("archived") else 0, int(thread["createdAt"]), + int(thread["updatedAt"]) if thread.get("updatedAt") is not None else None, thread.get("openaiCodeExecContainerId"), thread.get("anthropicCodeExecContainerId"), thread.get("forkedFromThreadId"), @@ -1084,6 +1108,7 @@ def update_chat_thread(id: str, patch: dict) -> Optional[dict]: "projectId": ("project_id", patch.get("projectId")), "archived": ("archived", 1 if patch.get("archived") else 0), "createdAt": ("created_at", patch.get("createdAt")), + "updatedAt": ("updated_at", patch.get("updatedAt")), "openaiCodeExecContainerId": ( "openai_code_exec_container_id", patch.get("openaiCodeExecContainerId"), @@ -1155,7 +1180,8 @@ def list_chat_threads( conn = get_connection() try: rows = conn.execute( - f"SELECT * FROM chat_threads {where} ORDER BY created_at DESC", + f"SELECT * FROM chat_threads {where} " + "ORDER BY COALESCE(updated_at, created_at) DESC, created_at DESC", values, ).fetchall() return [_chat_thread_from_row(row) for row in rows] @@ -1394,6 +1420,44 @@ def _raise_if_chat_message_thread_conflicts( ) +def _bump_chat_thread_updated_at( + conn: sqlite3.Connection, thread_id: str, message_created_at: int +) -> None: + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX(COALESCE(updated_at, created_at), ?) + WHERE id = ? + """, + (message_created_at, thread_id), + ) + + +def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str) -> None: + """Set updated_at from the remaining messages, floored at created_at. + + Unlike the ratchet-only bump, this can lower updated_at -- needed after + pruning, which may delete the thread's newest message. + """ + conn.execute( + """ + UPDATE chat_threads + SET updated_at = MAX( + COALESCE( + ( + SELECT MAX(m.created_at) FROM chat_messages m + WHERE m.thread_id = chat_threads.id + ), + created_at + ), + created_at + ) + WHERE id = ? + """, + (thread_id,), + ) + + def upsert_chat_message(message: dict) -> dict: conn = get_connection() try: @@ -1432,6 +1496,7 @@ def upsert_chat_message(message: dict) -> dict: int(message["createdAt"]), ), ) + _bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"])) conn.commit() return message except Exception: @@ -1484,6 +1549,12 @@ def sync_chat_messages( for m in messages ], ) + if prune_missing: + _recompute_chat_thread_updated_at(conn, thread_id) + elif messages: + _bump_chat_thread_updated_at( + conn, thread_id, max(int(m["createdAt"]) for m in messages) + ) conn.commit() return list_chat_messages(thread_id) except ChatMessageConflictError: diff --git a/studio/backend/tests/test_chat_history_storage.py b/studio/backend/tests/test_chat_history_storage.py index aa19df15fe..0239410734 100644 --- a/studio/backend/tests/test_chat_history_storage.py +++ b/studio/backend/tests/test_chat_history_storage.py @@ -4,6 +4,7 @@ import os import platform import shutil +import sqlite3 import threading import uuid from pathlib import Path @@ -11,6 +12,7 @@ from pathlib import Path import pytest from storage import studio_db +from utils.paths import studio_db_path def _reset_studio_db( @@ -108,6 +110,138 @@ def test_sync_chat_messages_upserts_without_pruning(tmp_path, monkeypatch): assert by_id["msg-2"]["content"] == [{"type": "text", "text": "updated text"}] +def test_chat_thread_updated_at_bumps_on_message_writes(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + assert thread["updatedAt"] == thread["createdAt"] + + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.upsert_chat_message(_message("msg-0", 1_600_000_000_000, "old")) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-2", 1_700_000_001_000, "newer")], + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + +def test_chat_thread_updated_at_recomputed_when_pruning(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + thread = studio_db.upsert_chat_thread(_thread()) + studio_db.sync_chat_messages( + "thread-1", + [ + _message("msg-1", 1_700_000_000_500, "older"), + _message("msg-2", 1_700_000_001_000, "newest"), + ], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_001_000 + + # Pruning the newest message must lower updated_at to the remaining one. + studio_db.sync_chat_messages( + "thread-1", + [_message("msg-1", 1_700_000_000_500, "older")], + prune_missing = True, + ) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + # Pruning every message falls back to created_at. + studio_db.sync_chat_messages("thread-1", [], prune_missing = True) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == thread["createdAt"] + + +def test_chat_thread_updated_at_survives_thread_resave(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + studio_db.upsert_chat_thread(_thread()) + studio_db.upsert_chat_message(_message("msg-1", 1_700_000_000_500, "hi")) + + studio_db.upsert_chat_thread(_thread()) + assert studio_db.get_chat_thread("thread-1")["updatedAt"] == 1_700_000_000_500 + + +def test_list_chat_threads_orders_by_last_activity(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + older = _thread("thread-old") + older["createdAt"] = 1_700_000_000_000 + newer = _thread("thread-new") + newer["createdAt"] = 1_700_000_100_000 + studio_db.upsert_chat_thread(older) + studio_db.upsert_chat_thread(newer) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-new", "thread-old"] + + studio_db.upsert_chat_message( + _message("msg-1", 1_700_000_200_000, "hi", thread_id = "thread-old") + ) + assert [t["id"] for t in studio_db.list_chat_threads()] == ["thread-old", "thread-new"] + + +def test_chat_threads_updated_at_migration_backfills_from_messages(tmp_path, monkeypatch): + _reset_studio_db(tmp_path, monkeypatch) + db_path = studio_db_path() + db_path.parent.mkdir(parents = True, exist_ok = True) + conn = sqlite3.connect(str(db_path)) + try: + conn.execute( + """ + CREATE TABLE chat_threads ( + id TEXT NOT NULL PRIMARY KEY, + title TEXT NOT NULL, + model_type TEXT NOT NULL, + model_id TEXT, + pair_id TEXT, + archived INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + """ + CREATE TABLE chat_messages ( + id TEXT NOT NULL PRIMARY KEY, + thread_id TEXT NOT NULL, + parent_id TEXT, + role TEXT NOT NULL, + content_json TEXT NOT NULL, + attachments_json TEXT, + metadata_json TEXT, + created_at INTEGER NOT NULL + ) + """ + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-with-msgs", "Old", "base", 1_700_000_000_000), + ) + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-empty", "Empty", "base", 1_700_000_050_000), + ) + # Fork-like thread: copied ancestor messages predate the thread itself. + conn.execute( + "INSERT INTO chat_threads (id, title, model_type, created_at) VALUES (?, ?, ?, ?)", + ("thread-fork", "Fork", "base", 1_700_000_100_000), + ) + conn.executemany( + "INSERT INTO chat_messages (id, thread_id, role, content_json, created_at) VALUES (?, ?, ?, ?, ?)", + [ + ("m1", "thread-with-msgs", "user", "[]", 1_700_000_001_000), + ("m2", "thread-with-msgs", "assistant", "[]", 1_700_000_002_000), + ("m3", "thread-fork", "user", "[]", 1_700_000_001_000), + ], + ) + conn.commit() + finally: + conn.close() + + assert studio_db.get_chat_thread("thread-with-msgs")["updatedAt"] == 1_700_000_002_000 + assert studio_db.get_chat_thread("thread-empty")["updatedAt"] == 1_700_000_050_000 + assert studio_db.get_chat_thread("thread-fork")["updatedAt"] == 1_700_000_100_000 + + def test_chat_projects_delete_cascades_threads_and_messages(tmp_path, monkeypatch): _reset_studio_db(tmp_path, monkeypatch) project = studio_db.upsert_chat_project(_project()) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index fb1a1fc9c7..2dd0d02515 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -359,7 +359,11 @@ export function AppSidebar() { const activeProjectId = isChatRoute ? ((search.project as string | undefined) ?? null) : null; - const { items: allChatItems } = useChatSidebarItems({ + const { + items: allChatItems, + archivedItems: archivedChatItems, + loaded: chatItemsLoaded, + } = useChatSidebarItems({ enabled: !isStudioRoute, requireMessages: false, }); @@ -1306,6 +1310,16 @@ export function AppSidebar() { renderChatSidebarItem(item, "recent"), )} + {/* "No chats yet" only when there is truly no history: + project-scoped and archived threads leave Recents empty + but still count as existing chats. */} + {chatItemsLoaded && + allChatItems.length === 0 && + archivedChatItems.length === 0 && ( +

+ {t("shell.navigation.noChatsYet")} +

+ )} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts index 55d7777e43..0a0df1139b 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts @@ -27,16 +27,21 @@ export interface SidebarItem { id: string; title: string; createdAt: number; + updatedAt: number; isFork?: boolean; projectId?: string | null; } +function lastActivityAt(thread: ThreadRecord): number { + return thread.updatedAt ?? thread.createdAt; +} + export function groupThreads( threads: ThreadRecord[], archived = false, ): SidebarItem[] { const items: SidebarItem[] = []; - const seenPairs = new Set(); + const pairItems = new Map(); for (const t of threads) { // Coerce archived to a boolean before comparing. Legacy threads (from the @@ -48,30 +53,35 @@ export function groupThreads( continue; } if (t.pairId) { - if (seenPairs.has(t.pairId)) { + const existing = pairItems.get(t.pairId); + if (existing) { + existing.updatedAt = Math.max(existing.updatedAt, lastActivityAt(t)); continue; } - seenPairs.add(t.pairId); - items.push({ + const item: SidebarItem = { type: "compare", id: t.pairId, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), projectId: t.projectId ?? null, - }); + }; + pairItems.set(t.pairId, item); + items.push(item); } else if (!t.pairId) { items.push({ type: "single", id: t.id, title: t.title, createdAt: t.createdAt, + updatedAt: lastActivityAt(t), isFork: Boolean(t.forkedFromThreadId), projectId: t.projectId ?? null, }); } } - return items.sort((a, b) => b.createdAt - a.createdAt); + return items.sort((a, b) => b.updatedAt - a.updatedAt); } // Streaming fires CHAT_HISTORY_UPDATED_EVENT per chunk. Debounce so each quiet @@ -84,6 +94,7 @@ export function useChatSidebarItems(options?: { requireMessages?: boolean; }) { const [allThreads, setAllThreads] = useState([]); + const [loaded, setLoaded] = useState(false); const enabled = options?.enabled ?? true; const requireMessages = options?.requireMessages ?? true; @@ -111,6 +122,7 @@ export function useChatSidebarItems(options?: { // were in flight, or if the effect was torn down. if (cancelled || seq !== requestSeq) return; setAllThreads(threads); + setLoaded(true); } catch (error) { if (isExpectedBackgroundChatStorageError(error)) { return; @@ -144,7 +156,7 @@ export function useChatSidebarItems(options?: { const archivedItems = groupThreads(allThreads ?? [], true); const canCompare = useChatRuntimeStore((s) => Boolean(s.params.checkpoint)); - return { items, archivedItems, canCompare }; + return { items, archivedItems, canCompare, loaded }; } function cancelIfRunning(threadId: string): void { diff --git a/studio/frontend/src/features/chat/types.ts b/studio/frontend/src/features/chat/types.ts index 3fe69ccc26..111510d925 100644 --- a/studio/frontend/src/features/chat/types.ts +++ b/studio/frontend/src/features/chat/types.ts @@ -36,6 +36,7 @@ export interface ThreadRecord { projectId?: string | null; archived: boolean; createdAt: number; + updatedAt?: number; /** * OpenAI shell tool container id from a prior response. When set, the * next turn reuses it via `environment.type="container_reference"` so diff --git a/studio/frontend/src/features/chat/utils/chat-history-storage.ts b/studio/frontend/src/features/chat/utils/chat-history-storage.ts index 4294887df0..00df3657b1 100644 --- a/studio/frontend/src/features/chat/utils/chat-history-storage.ts +++ b/studio/frontend/src/features/chat/utils/chat-history-storage.ts @@ -590,7 +590,10 @@ export async function listStoredChatThreads( } return Array.from(byId.values()) .filter((thread) => matchesThreadListArgs(thread, args)) - .sort((a, b) => b.createdAt - a.createdAt); + .sort( + (a, b) => + (b.updatedAt ?? b.createdAt) - (a.updatedAt ?? a.createdAt), + ); } export async function listStoredChatThreadsWithMessages( diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index 3d73ad4343..b67fd5ca1d 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -41,6 +41,7 @@ export const en = { recipes: "Recipes", export: "Export", recents: "Recents", + noChatsYet: "No chats yet", settings: "Settings", api: "API", lightMode: "Light Mode", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index 5fd31dc7cb..f6dc265fc7 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -41,6 +41,7 @@ export const zhCN = { recipes: "配方", export: "导出", recents: "最近", + noChatsYet: "暂无对话", settings: "设置", api: "API", lightMode: "浅色模式", From 93c9d6d0dd0bfd48d5766ac952c3a880a47b3727 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 7 Jul 2026 15:13:53 -0300 Subject: [PATCH 003/402] Studio: render \[ \] and \( \) LaTeX delimiters in chat (#6914) --- studio/frontend/src/lib/latex.ts | 169 ++++++++++++++++++++++++++++--- 1 file changed, 155 insertions(+), 14 deletions(-) diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 1a4ebdace8..86a9634048 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -1,8 +1,13 @@ // Adapted from LibreChat's latex.ts // https://github.com/danny-avila/LibreChat/blob/main/client/src/utils/latex.ts // -// Escapes currency dollar signs so they are not misinterpreted as LaTeX math -// delimiters when singleDollarTextMath is enabled. +// Two jobs, in order: +// 1. Convert LaTeX bracket delimiters (`\[...\]`, `\(...\)`) into the dollar +// forms remark-math understands (`$$...$$`, `$...$`). remark-math only +// tokenizes dollar delimiters, so models that emit `\[...\]` / `\(...\)` +// would otherwise render as literal text. +// 2. Escape currency dollar signs so they are not misinterpreted as LaTeX +// math delimiters when singleDollarTextMath is enabled. /** * Matches a single $ followed by a number pattern (currency), e.g.: @@ -15,14 +20,14 @@ const CURRENCY_REGEX = /(? { const regions: Array<[number, number]> = []; - // Fenced code blocks: ```...``` - const fencedRe = /```[\s\S]*?```/g; + // Fenced code blocks: ```...``` and ~~~...~~~ (both are code in GFM) + const fencedRe = /```[\s\S]*?```|~~~[\s\S]*?~~~/g; let match: RegExpExecArray | null; while ((match = fencedRe.exec(content)) !== null) { regions.push([match.index, match.index + match[0].length]); @@ -51,9 +56,38 @@ function findCodeBlockRegions(content: string): Array<[number, number]> { } /** - * Binary search to check if a position falls inside any code region. + * Match an inline link/image `[text](DEST)`, capturing the destination as group 1 + * with the `d` flag so its span is read straight from `match.indices` (the text + * can contain an escaped `\](`, so a string search for the separator is unsafe). + * The text disallows unescaped `]`; the destination allows escapes and one level + * of balanced parens. */ -function isInCodeBlock( +const LINK_DEST_RE = + /!?\[(?:\\.|[^\]\\])*?\]\(((?:\\.|[^()\\]|\([^()]*\))*)\)/gd; + +/** + * Find the destination spans of inline links/images, so a `\(...\)` written with + * escaped parens inside a URL isn't rewritten as math (which would break the + * link). Only the destination is returned, not the link text, so math in the + * visible text still converts. Sorted, non-overlapping (matches are disjoint). + */ +function findLinkDestinationRegions(content: string): Array<[number, number]> { + if (!content.includes("](")) return []; + const regions: Array<[number, number]> = []; + let match: RegExpExecArray | null; + LINK_DEST_RE.lastIndex = 0; + while ((match = LINK_DEST_RE.exec(content)) !== null) { + // `indices` is present (the `d` flag); group 1 spans the destination. + regions.push(match.indices![1]); + } + return regions; +} + +/** + * Binary search to check if a position falls inside any region. Regions must be + * sorted by start and non-overlapping. + */ +function isInRegion( position: number, regions: Array<[number, number]>, ): boolean { @@ -174,9 +208,109 @@ function hasInlineMathCloser(content: string, offset: number): boolean { } /** - * Preprocess a markdown string to escape currency dollar signs so they are not - * parsed as LaTeX math delimiters. + * Matches a `\[...\]` (display) or `\(...\)` (inline) LaTeX span. Non-greedy so + * the first closer wins; dotall so display spans can wrap lines. `(? block `$$...$$` and `\(...\)` -> inline `$...$` so + * remark-math can tokenize them. Bodies are trimmed: remark-math won't open an + * inline span on `$ ` (a `$` followed by whitespace), and display fences must + * sit on their own line to render as a centered block (not inline math), so + * `\[...\]` becomes `\n$$\n...\n$$\n`. * + * Spans inside code blocks/spans are left intact (a code sample showing `\(x\)` + * must not be rewritten). + * + * A space is inserted between a converted span and a following `$` so their + * delimiters can't fuse (`\(a\)\(b\)` -> `$a$$b$` would mis-tokenize into one + * broken span). A preceding currency (`$5\(x\)`) is instead broken later by the + * currency escape pass. + * + * Returns the rewritten text and the `[start, end)` ranges (in the rewritten + * string) of every span it produced, so the currency pass can skip them. + */ +function convertLatexDelimiters(content: string): { + text: string; + mathRegions: Array<[number, number]>; +} { + if (!content.includes("\\[") && !content.includes("\\(")) { + return { text: content, mathRegions: [] }; + } + + const codeRegions = findCodeBlockRegions(content); + const linkRegions = findLinkDestinationRegions(content); + const inSkipZone = (pos: number) => + isInRegion(pos, codeRegions) || isInRegion(pos, linkRegions); + // Pushed in ascending, non-overlapping order (offset only grows), so this + // stays valid for isInRegion's binary search without a sort. + const mathRegions: Array<[number, number]> = []; + // Accumulate into an array, not a string: reading the last char off a growing + // `+=` accumulator flattens its rope every append (O(n^2) over many spans, on + // the per-frame streaming path), so track the tail char and length instead. + const parts: string[] = []; + let offset = 0; + let lastChar = ""; + let last = 0; + // Append a chunk, separating a trailing `$` from a leading `$` so two spans + // can't fuse. Returns where the chunk landed (after any inserted space). + const append = (chunk: string): number => { + if (!chunk) return offset; + if (lastChar === "$" && chunk.startsWith("$")) { + parts.push(" "); + offset += 1; + } + const start = offset; + parts.push(chunk); + offset += chunk.length; + lastChar = chunk[chunk.length - 1]; + return start; + }; + let match: RegExpExecArray | null; + LATEX_DELIM_RE.lastIndex = 0; + while ((match = LATEX_DELIM_RE.exec(content)) !== null) { + const matchEnd = match.index + match[0].length; + // Skip if either delimiter is inside code or a link destination: an opener + // outside such a zone must not consume a closer inside one and rewrite + // across the boundary. Resume right after this opener (not past the whole + // match) so a valid span that this match spanned across (a stray code `\(` + // paired with a real closer) is still found on the next pass, not swallowed. + if (inSkipZone(match.index) || inSkipZone(matchEnd - 1)) { + LATEX_DELIM_RE.lastIndex = match.index + 1; + continue; + } + const isDisplay = match[1] !== undefined; + const body = (isDisplay ? match[1] : match[2]).trim(); + // Leave an empty span (`\(\)`) literal; a bare `$$` would open a stray + // display block that swallows following text. + if (!body) { + continue; + } + append(content.slice(last, match.index)); + const wrapped = isDisplay ? `\n$$\n${body}\n$$\n` : `$${body}$`; + const start = append(wrapped); + mathRegions.push([start, offset]); + last = matchEnd; + } + append(content.slice(last)); + return { text: parts.join(""), mathRegions }; +} + +/** + * Preprocess a markdown string so LaTeX renders: convert bracket delimiters to + * dollar forms, then escape currency dollar signs so they are not parsed as + * math delimiters. + * + * - `\[E = mc^2\]` becomes a `$$` display block on its own lines (display math) + * - `\(\alpha\)` becomes `$\alpha$` (inline math) + * - `\(x\)` in a code span is untouched * - `$5` alone becomes `\$5` (currency, not math) * - `$\alpha$` is untouched (real LaTeX) * - `$30^\circ$` is untouched (LaTeX whose body starts with a digit) @@ -185,15 +319,22 @@ function hasInlineMathCloser(content: string, offset: number): boolean { * - Currency inside code blocks/spans is untouched */ export function preprocessLaTeX(content: string): string { - if (!content.includes("$")) return content; + const { text, mathRegions } = convertLatexDelimiters(content); - const codeRegions = findCodeBlockRegions(content); + if (!text.includes("$")) return text; - return content.replace(CURRENCY_REGEX, (match, offset) => { - if (isInCodeBlock(offset, codeRegions)) { + const codeRegions = findCodeBlockRegions(text); + + return text.replace(CURRENCY_REGEX, (match, offset) => { + if (isInRegion(offset, codeRegions)) { return match; } - if (hasInlineMathCloser(content, offset)) { + // Skip the spans we just created from `\(...\)` so a numeric body like + // `$5$` isn't re-escaped back to literal `\$5$`. + if (isInRegion(offset, mathRegions)) { + return match; + } + if (hasInlineMathCloser(text, offset)) { return match; } return "\\" + match; From 304b8eca7ae8a7bd743850a248308931c062dab9 Mon Sep 17 00:00:00 2001 From: Ayushman <139611211+InfoSage05@users.noreply.github.com> Date: Tue, 7 Jul 2026 23:58:18 +0530 Subject: [PATCH 004/402] fix: match qwen3-thinking double-newline in train_on_responses_only response pattern (#6926) * fix: match qwen3-thinking chat template double-newline in response pattern The Qwen3-thinking chat template generates `\n\n` (double newline) after the think tag, but `train_on_responses_only` was looking for `\n` (single newline). `\n\n` is token 271 while `\n` is token 198 -- different tokens, so the pattern match in `train_on_responses_only` fails, masking ALL tokens and dropping 100% of training samples. Update the response pattern from `\n` to `\n\n` to match what the actual qwen3-thinking template generates. Fixes #6919 * fix qwen3 thinking response marker --------- Co-authored-by: Ayushman Paul Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> --- studio/backend/utils/datasets/model_mappings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/studio/backend/utils/datasets/model_mappings.py b/studio/backend/utils/datasets/model_mappings.py index 463d26a692..9d2c983aed 100644 --- a/studio/backend/utils/datasets/model_mappings.py +++ b/studio/backend/utils/datasets/model_mappings.py @@ -487,7 +487,7 @@ TEMPLATE_TO_RESPONSES_MAPPER = { }, "qwen3-thinking": { "instruction": "<|im_start|>user\n", - "response": "<|im_start|>assistant\n\n", + "response": "<|im_start|>assistant\n", }, "qwen3": { "instruction": "<|im_start|>user\n", From a9db53e189f2c23586bfc4a6f472448afc4807ef Mon Sep 17 00:00:00 2001 From: oobabooga Date: Tue, 7 Jul 2026 19:50:40 -0300 Subject: [PATCH 005/402] Studio: stream reasoning tokens in the tool-loop generator (fixes DeepSeek thinking not streaming with a pill on) (#6947) --- .../core/inference/anthropic_compat.py | 28 +++ studio/backend/core/inference/llama_cpp.py | 54 +++- .../backend/tests/test_anthropic_messages.py | 43 ++++ .../backend/tests/test_llama_cpp_tool_loop.py | 231 +++++++++++++++++- 4 files changed, 338 insertions(+), 18 deletions(-) diff --git a/studio/backend/core/inference/anthropic_compat.py b/studio/backend/core/inference/anthropic_compat.py index 7b572a28ff..3c7a4cb182 100644 --- a/studio/backend/core/inference/anthropic_compat.py +++ b/studio/backend/core/inference/anthropic_compat.py @@ -258,6 +258,10 @@ class AnthropicStreamEmitter: self._open_tool_use_id: Optional[str] = None self._open_tool_args_sent: bool = False self._prev_text: str = "" + # Net minus in the text emitted to the client. Tracked + # from emitted deltas (not _prev_text, which a final bare shrink clobbers) + # so an unclosed reasoning-only block can be balanced before close. + self._open_think_tags: int = 0 self._usage: dict = {} def start( @@ -317,6 +321,7 @@ class AnthropicStreamEmitter: """Close any open block and emit message_delta + message_stop.""" events = [] if self._text_block_open or self._open_tool_call_id is not None: + events.extend(self._close_open_think()) events.append(self._close_block()) self._open_tool_call_id = None self._open_tool_use_id = None @@ -344,12 +349,33 @@ class AnthropicStreamEmitter: ) return events + def _close_open_think(self) -> list[str]: + """Emit a ```` delta when the streamed text left a ```` + open. This emitter diffs cumulative snapshots and drops the generator's + final bare shrink, so a reasoning-only reply would otherwise end on an + unclosed tag. Mirrors the chat route's reasoning extractor, which closes + the block on finish; balances the block before it is closed.""" + if not self._text_block_open or self._open_think_tags <= 0: + return [] + self._open_think_tags = 0 + return [ + build_anthropic_sse_event( + "content_block_delta", + { + "type": "content_block_delta", + "index": self.block_index, + "delta": {"type": "text_delta", "text": ""}, + }, + ) + ] + def _handle_content(self, event: dict) -> list[str]: cumulative = event.get("text", "") new_text = cumulative[len(self._prev_text) :] self._prev_text = cumulative if not new_text: return [] + self._open_think_tags += new_text.count("") - new_text.count("") if not self._text_block_open: events = self._open_text_block() else: @@ -374,6 +400,7 @@ class AnthropicStreamEmitter: events = [] if self._text_block_open: + events.extend(self._close_open_think()) events.append(self._close_block()) # Defensive: close a stale open tool_use block before starting another. elif self._open_tool_call_id is not None: @@ -452,6 +479,7 @@ class AnthropicStreamEmitter: events.extend(self._open_text_block()) # Reset text tracking for the next synthesis turn self._prev_text = "" + self._open_think_tags = 0 return events def _open_text_block(self) -> list[str]: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 11a4ebb3ec..757467a008 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -8603,13 +8603,31 @@ class LlamaCppBackend: } def _flush_reasoning_and_buffer(): - """Append buffered reasoning (as a block) then the held + """Close a live-streamed block (or emit the buffered reasoning + as one block if it never streamed), then append the held content_buffer to the cumulative display text.""" - nonlocal cumulative_display - if reasoning_accum: + nonlocal cumulative_display, in_thinking + if in_thinking: + cumulative_display += "" + in_thinking = False + elif reasoning_accum: cumulative_display += "" + reasoning_accum + "" cumulative_display += content_buffer + def _close_streamed_think() -> bool: + """Close a live-streamed before a tool call drains, so + consumers without a reasoning extractor (Anthropic) get a balanced + block. Returns True when the caller should yield the result.""" + nonlocal cumulative_display, in_thinking, _last_emitted + if not in_thinking: + return False + cumulative_display += "" + in_thinking = False + if len(cumulative_display) > len(_last_emitted) and not _suppress_visible_output: + _last_emitted = cumulative_display + return True + return False + def _looks_like_enabled_bare_json(text: str, enabled_tool_names: set) -> bool: """True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False.""" probe = strip_llama3_leading_sentinels(text.lstrip()) @@ -8797,6 +8815,10 @@ class LlamaCppBackend: # the structured tool call. has_structured_tc = True detect_state = _S_DRAINING + # Close the reasoning prefix before the tool card + # (mirrors the is_match path). + if _close_streamed_think(): + yield {"type": "content", "text": cumulative_display} for tc_d in tc_deltas: idx = tc_d.get("index", 0) if idx not in tool_calls_acc: @@ -8882,17 +8904,17 @@ class LlamaCppBackend: continue # ── Reasoning tokens ── - # Yield only in STREAMING. In BUFFERING and - # DRAINING, accumulate silently so we don't - # corrupt the consumer's prev_text tracker - # (routes/inference.py never resets it - # between tool iterations). + # Stream live except while DRAINING: reasoning is + # orthogonal to tool detection (content_buffer + # only), and the route resets prev_text on + # tool_start, so the block stays a + # monotonic prefix like the no-tool path. reasoning = delta.get("reasoning_content", "") if reasoning: if _reasoning_started_at is None: _reasoning_started_at = time.monotonic() reasoning_accum += reasoning - if detect_state == _S_STREAMING: + if detect_state != _S_DRAINING: if not in_thinking: cumulative_display += "" in_thinking = True @@ -9020,9 +9042,15 @@ class LlamaCppBackend: _hold_buffer = True if _drain_silently: - # No visible prefix -- the buffered text IS - # the call; drain without yielding it. + # The buffered content IS the call; drain it + # without yielding. A live prefix is + # separate from it -- close that. detect_state = _S_DRAINING + if _close_streamed_think(): + yield { + "type": "content", + "text": cumulative_display, + } elif is_match: # Tool signal -- flush any visible # prefix before DRAINING so the @@ -9115,7 +9143,9 @@ class LlamaCppBackend: ), } elif reasoning_accum and not has_content_tokens: - # Reasoning-only reply: show it as plain text. + # Reasoning-only reply: show it as the main response, + # not a thinking block (mirrors the no-tool path; the + # route's extractor closes the streamed ). if _reasoning_started_at is not None and not _reasoning_summary_emitted: _reasoning_summary_emitted = True yield _reasoning_summary_event(_reasoning_started_at) diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 170b456eac..0c6550a3bb 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -52,6 +52,49 @@ from io import BytesIO as _BytesIO from types import SimpleNamespace +def _emitter_client_text(events: list[str]) -> str: + """Concatenate the text_delta payloads an SSE event list carries.""" + text = "" + for line in events: + for raw in line.split("\n"): + raw = raw.strip() + if not raw.startswith("data: "): + continue + data = json.loads(raw[len("data: ") :]) + delta = data.get("delta", {}) + if delta.get("type") == "text_delta": + text += delta.get("text", "") + return text + + +def test_anthropic_emitter_closes_reasoning_only_think_block(): + # A reasoning-only reply streams X live then shrinks to bare X at EOF. + # This emitter diffs cumulative snapshots and drops the shrink, so without a + # closing pass the client text would end on an unclosed . finish() + # must balance it. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "The capital"}) + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + # The generator's final bare-text shrink (dropped by the cumulative diff). + events += emitter.feed({"type": "content", "text": "The capital of France is Paris."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "The capital of France is Paris." + + +def test_anthropic_emitter_does_not_double_close_balanced_think(): + # A reasoning-then-answer reply already closes its own ; the balancer + # must not append a second one. + emitter = AnthropicStreamEmitter() + events = emitter.start("msg_1", "m") + events += emitter.feed({"type": "content", "text": "Thinking."}) + events += emitter.feed({"type": "content", "text": "Thinking.Answer."}) + events += emitter.finish() + + assert _emitter_client_text(events) == "Thinking.Answer." + + def test_streamed_anthropic_tool_use_records_api_monitor_reply(monkeypatch): import routes.inference as inf_mod diff --git a/studio/backend/tests/test_llama_cpp_tool_loop.py b/studio/backend/tests/test_llama_cpp_tool_loop.py index fb1b0e52b7..afac1f5249 100644 --- a/studio/backend/tests/test_llama_cpp_tool_loop.py +++ b/studio/backend/tests/test_llama_cpp_tool_loop.py @@ -221,7 +221,7 @@ def test_structured_tool_call_after_visible_preface_is_executed(monkeypatch): assert assistant_messages[-1]["tool_calls"][0]["function"]["name"] == "render_html" -def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): +def test_streamed_reasoning_answer_emits_backend_summary(monkeypatch): stream = [ _sse({"reasoning_content": "I am thinking."}), _sse({"reasoning_content": " Still thinking."}), @@ -240,17 +240,236 @@ def test_buffered_reasoning_answer_emits_backend_summary(monkeypatch): ) ) + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streams live during BUFFERING instead of arriving as one block: + # each reasoning delta is emitted immediately, wrapped in . + assert content_texts[0] == "I am thinking." + assert content_texts[1] == "I am thinking. Still thinking." + # The final event closes the block and appends the answer. + assert content_texts[-1] == "I am thinking. Still thinking.Final answer." + summary_index = next( i for i, event in enumerate(events) if event["type"] == "reasoning_summary" ) - content_index = next(i for i, event in enumerate(events) if event["type"] == "content") - assert summary_index < content_index + final_content_index = max(i for i, event in enumerate(events) if event["type"] == "content") + assert summary_index < final_content_index assert events[summary_index]["duration_ms"] == 62000 - assert ( - events[content_index]["text"] - == "I am thinking. Still thinking.Final answer." + + +def test_reasoning_streams_incrementally_with_tools(monkeypatch): + # Regression (DeepSeek "thinking doesn't stream"): with a tool/pill active the + # tool-loop generator must stream reasoning token-by-token like the no-tool + # path, not accumulate it and dump one buffered block. + stream = [ + _sse({"reasoning_content": "Step one."}), + _sse({"reasoning_content": " Step two."}), + _sse({"reasoning_content": " Step three."}), + _sse({"content": "Done."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "think then answer"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) ) + reasoning_stage = [ + e["text"] + for e in events + if e["type"] == "content" + and e["text"].startswith("") + and "" not in e["text"] + ] + # One live emission per reasoning delta -- not a single dump. + assert reasoning_stage == [ + "Step one.", + "Step one. Step two.", + "Step one. Step two. Step three.", + ] + final = [e["text"] for e in events if e["type"] == "content"][-1] + assert final == "Step one. Step two. Step three.Done." + + +def test_reasoning_only_reply_matches_no_tool_path_with_tools(monkeypatch): + # A reasoning-only turn (whole answer in reasoning_content, no content, no + # tool) with a tool active streams the reasoning live, then resolves to the + # bare reasoning text -- identical to the no-tool generate_chat_completion + # path -- so the non-streaming drain still returns it as `content`, not an + # empty answer. + stream = [ + _sse({"reasoning_content": "The capital of France is Paris."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 5.0, 5.0]) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "just think"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + content_texts = [e["text"] for e in events if e["type"] == "content"] + # Reasoning streamed live during BUFFERING (the fix). + assert content_texts[0] == "The capital of France is Paris." + # Resolves to bare reasoning, matching the no-tool sibling. + assert content_texts[-1] == "The capital of France is Paris." + + +def test_reasoning_before_structured_tool_closes_think_block(monkeypatch): + # Regression: reasoning streamed live during BUFFERING must be closed with + # before a structured tool_call drains, so consumers without a + # reasoning extractor (Anthropic /v1/messages) never receive an unclosed + # . Mirrors the is_match (XML tool signal) path. + tool_stream = [ + _sse({"reasoning_content": "Let me search."}), + *_structured_tool_call("web_search", {"query": "weather"}, "call_1"), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + # Reasoning streamed live, then closed before the tool -- balanced block. + assert content_before_tool[0] == "Let me search." + assert content_before_tool[-1] == "Let me search." + + +def _replay_route_reasoning_extractor(cumulatives: list[str]) -> tuple[str, str]: + """Replay the route's cumulative suffix-diff + reasoning extractor (the + shared core of routes/inference.py gguf_stream_chunks and the tool-loop + consumer) over content snapshots. Returns (visible, reasoning).""" + from routes.inference import _ResponsesReasoningExtractor + + extractor = _ResponsesReasoningExtractor(parse_think_markers = True) + prev_text = "" + visible: list[str] = [] + reasoning: list[str] = [] + for cumulative in cumulatives: + new_text = cumulative[len(prev_text) :] + prev_text = cumulative + if not new_text: + continue + reasoning_delta, visible_delta = extractor.feed(new_text) + if reasoning_delta: + reasoning.append(reasoning_delta) + if visible_delta: + visible.append(visible_delta) + final_reasoning, final_visible = extractor.finish() + if final_reasoning: + reasoning.append(final_reasoning) + if final_visible: + visible.append(final_visible) + return "".join(visible), "".join(reasoning) + + +def test_reasoning_only_route_output_matches_no_tool_path(monkeypatch): + # Parity contract: a reasoning-only reply must reach the client identically + # whether tools are on or off. Both generators stream live then + # resolve to the bare reasoning text; the route's suffix-diff + extractor + # must therefore produce the same (visible, reasoning) split for both. + stream = [ + _sse({"reasoning_content": "The capital"}), + _sse({"reasoning_content": " of France is Paris."}), + _done(), + ] + + tool_backend = _make_backend(monkeypatch, [list(stream)], []) + _patch_monotonic(monkeypatch, [1.0, 2.0, 2.0]) + tool_cumulatives = [ + e["text"] + for e in tool_backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "capital of France?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + if e.get("type") == "content" + ] + + no_tool_backend = _make_backend(monkeypatch, [list(stream)], []) + no_tool_cumulatives = [ + y + for y in no_tool_backend.generate_chat_completion( + messages = [{"role": "user", "content": "capital of France?"}], + ) + if isinstance(y, str) + ] + + # Both paths stream the reasoning live with the same leading shape. (Raw + # yield lists aren't compared verbatim: the tool path emits a pre-existing + # duplicate trailing event that the route's suffix-diff dedupes.) + assert tool_cumulatives[:3] == no_tool_cumulatives[:3] + # The contract that matters: identical route-level output. + tool_out = _replay_route_reasoning_extractor(tool_cumulatives) + no_tool_out = _replay_route_reasoning_extractor(no_tool_cumulatives) + assert tool_out == no_tool_out + # Pin the shared contract so a change to either path shows up here. + _visible, reasoning = tool_out + assert reasoning == "The capital of France is Paris." + + +def test_reasoning_before_bare_json_tool_closes_think_block(monkeypatch): + # _drain_silently sibling of the structured-tool close: a bare-JSON tool call + # with a live reasoning prefix must also close before draining, and + # must never leak the drained call text as content. + tool_stream = [ + _sse({"reasoning_content": "Searching now."}), + _sse({"content": '{"name":"web_search","arguments":{"query":"weather"}}'}), + _done(), + ] + final_stream = [ + _sse({"content": "It is sunny."}), + _done(), + ] + payloads: list[dict] = [] + backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads) + _patch_monotonic(monkeypatch, [1.0, 2.0, 3.0, 4.0, 4.0]) + + monkeypatch.setattr( + "core.inference.tools.execute_tool", lambda name, arguments, **_kwargs: "sunny" + ) + + events = list( + backend.generate_chat_completion_with_tools( + messages = [{"role": "user", "content": "weather?"}], + tools = [{"type": "function", "function": {"name": "web_search"}}], + max_tool_iterations = 1, + ) + ) + + tool_start_index = next(i for i, e in enumerate(events) if e["type"] == "tool_start") + content_before_tool = [e["text"] for e in events[:tool_start_index] if e["type"] == "content"] + assert content_before_tool[0] == "Searching now." + assert content_before_tool[-1] == "Searching now." + # The bare-JSON call text was drained, never surfaced as content. + assert not any('"name"' in t for t in content_before_tool) + def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch): tool_stream = [ From 01b8085dc2e8fae5d99ee7d236d58b706988773c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Tue, 7 Jul 2026 17:10:01 -0700 Subject: [PATCH 006/402] Create ossf.yml (#6952) --- .github/workflows/ossf.yml | 78 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 78 insertions(+) create mode 100644 .github/workflows/ossf.yml diff --git a/.github/workflows/ossf.yml b/.github/workflows/ossf.yml new file mode 100644 index 0000000000..f9a270540f --- /dev/null +++ b/.github/workflows/ossf.yml @@ -0,0 +1,78 @@ +# This workflow uses actions that are not certified by GitHub. They are provided +# by a third-party and are governed by separate terms of service, privacy +# policy, and support documentation. + +name: Scorecard supply-chain security +on: + # For Branch-Protection check. Only the default branch is supported. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection + branch_protection_rule: + # To guarantee Maintained check is occasionally updated. See + # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained + schedule: + - cron: '21 20 * * 0' + push: + branches: [ "main" ] + +# Declare default permissions as read only. +permissions: read-all + +jobs: + analysis: + name: Scorecard analysis + runs-on: ubuntu-latest + # `publish_results: true` only works when run from the default branch. conditional can be removed if disabled. + if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request' + permissions: + # Needed to upload the results to code-scanning dashboard. + security-events: write + # Needed to publish results and get a badge (see publish_results below). + id-token: write + # Uncomment the permissions below if installing in a private repository. + # contents: read + # actions: read + + steps: + - name: "Checkout code" + uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2 + with: + persist-credentials: false + + - name: "Run analysis" + uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1 + with: + results_file: results.sarif + results_format: sarif + # (Optional) "write" PAT token. Uncomment the `repo_token` line below if: + # - you want to enable the Branch-Protection check on a *public* repository, or + # - you are installing Scorecard on a *private* repository + # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional. + # repo_token: ${{ secrets.SCORECARD_TOKEN }} + + # Public repositories: + # - Publish results to OpenSSF REST API for easy access by consumers + # - Allows the repository to include the Scorecard badge. + # - See https://github.com/ossf/scorecard-action#publishing-results. + # For private repositories: + # - `publish_results` will always be set to `false`, regardless + # of the value entered here. + publish_results: true + + # (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore + # file_mode: git + + # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF + # format to the repository Actions tab. + - name: "Upload artifact" + uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1 + with: + name: SARIF file + path: results.sarif + retention-days: 5 + + # Upload the results to GitHub's code scanning dashboard (optional). + # Commenting out will disable upload of results to your repo's Code Scanning dashboard + - name: "Upload to code-scanning" + uses: github/codeql-action/upload-sarif@v3 + with: + sarif_file: results.sarif From 49d1fb38633b2e5b640f7034a3e69734a16396ac Mon Sep 17 00:00:00 2001 From: Wasim Yousef Said Date: Wed, 8 Jul 2026 03:08:07 +0200 Subject: [PATCH 007/402] Speed up Studio startup path (#6899) * Speed up Studio startup path * Studio: recheck managed binary executability on preflight cache hit and ignore stale unauthenticated platform fetches Preflight: a matching capability cache fingerprint no longer skips the runnability check when the managed binary's executable bit was cleared (size and mtime unchanged, since chmod bumps ctime not mtime). The cache fast path now confirms the binary is still executable, otherwise it falls back to the CLI help probe so preflight reports Stale and can repair, instead of returning Ready and failing later at backend start. Adds a regression test. Frontend: now that first render is no longer gated on fetchDeviceType, the initial unauthenticated health call can resolve after an authenticated platform fetch. Guard the store so a late unauthenticated or failed non-forced response cannot overwrite an already authoritative device type, tunnel URL, or secure flag. Forced refreshes and the first unauthenticated load are unaffected. * Studio: use access(X_OK) for the preflight cache executability guard A mode bitmask treats any execute bit as launchable, but the executable bits can be set only for another owner or group, or be denied by an ACL, so the current user could still hit PermissionDenied at launch and the cached fast path would wrongly return Ready. access(X_OK) checks real executability for the calling user, so an ownership or permission change correctly falls back to the CLI help probe and the Stale repair path. * Studio: ignore any stale non-forced platform fetch once authoritative Extend the platform store guard so a non-forced health response never overwrites an already authoritative result, not only unauthenticated ones. With a saved token the post-render non-forced request can be authenticated but older than a later forced refresh that already picked up the tunnel URL and secure flag; if that earlier request resolves last it would null those fields. Now any non-forced response is dropped once the store holds a server-reported platform. Forced refreshes and the first authoritative write are unaffected. * Studio: run the managed CLI help probe before trusting the preflight cache Restore running the managed CLI help probe before returning Ready from the desktop capability cache, so a managed install whose venv interpreter or a runtime dependency is broken (while path, size, mtime, and markers are unchanged) is reported Stale for repair rather than proceeding to a backend start that cannot spawn. The capability cache still skips the heavier desktop-capabilities probe on a hit, so a warm cache runs one probe instead of two. Removes the executable-access shortcut, which the help probe now subsumes. --------- Co-authored-by: Daniel Han --- studio/backend/core/inference/orchestrator.py | 8 +- ...t_inference_default_models_non_blocking.py | 42 ++++++ .../frontend/src/components/app-sidebar.tsx | 45 ++++-- studio/frontend/src/config/env.ts | 26 +++- .../frontend/src/features/chat/chat-page.tsx | 43 +++++- .../chat/hooks/use-chat-model-runtime.ts | 12 +- .../src/features/chat/runtime-provider.tsx | 12 +- studio/frontend/src/main.tsx | 16 +-- studio/src-tauri/src/preflight.rs | 132 +++++++++++++++++- studio/src-tauri/src/preflight/managed.rs | 16 +++ 10 files changed, 308 insertions(+), 44 deletions(-) create mode 100644 studio/backend/tests/test_inference_default_models_non_blocking.py diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index 19d2230278..cf5d24c367 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -108,13 +108,11 @@ class InferenceOrchestrator: @property def default_models(self) -> list[str]: - # Wait up to 5s for background HF fetch - self._top_models_ready.wait(timeout = 5) top_gguf = self._top_gguf_cache or [] top_hub = self._top_hub_cache or [] - # Curated static defaults first, then HF download-ranked to backfill. - # Send extras so the frontend keeps 4 per category after removing - # downloaded ones. + # Never wait for the remote Hugging Face ranking during startup. Chat's + # first /api/models/list needs curated defaults immediately; the + # background fetch backfills extra choices on later calls. result: list[str] = [] seen: set[str] = set() for m in self._static_models + top_gguf + top_hub: diff --git a/studio/backend/tests/test_inference_default_models_non_blocking.py b/studio/backend/tests/test_inference_default_models_non_blocking.py new file mode 100644 index 0000000000..83a8e7bbfb --- /dev/null +++ b/studio/backend/tests/test_inference_default_models_non_blocking.py @@ -0,0 +1,42 @@ +# 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 Chat model metadata must not block on remote Hugging Face discovery.""" + +from __future__ import annotations + +import sys +import time +from pathlib import Path + +_BACKEND = Path(__file__).resolve().parent.parent +if str(_BACKEND) not in sys.path: + sys.path.insert(0, str(_BACKEND)) + +from core.inference.orchestrator import InferenceOrchestrator # noqa: E402 + + +def test_default_models_returns_static_defaults_before_top_fetch(monkeypatch): + sleep_seconds = 2.0 + + def _slow_fetch(self: InferenceOrchestrator) -> None: + time.sleep(sleep_seconds) + self._top_gguf_cache = ["unsloth/slow-GGUF"] + self._top_models_ready.set() + + monkeypatch.setattr(InferenceOrchestrator, "_fetch_top_models", _slow_fetch) + + orchestrator = InferenceOrchestrator() + started = time.monotonic() + defaults = orchestrator.default_models + elapsed = time.monotonic() - started + + assert elapsed < 0.5, f"default_models blocked for {elapsed:.2f}s" + assert defaults == orchestrator._static_models + assert "unsloth/slow-GGUF" not in defaults + + deadline = time.monotonic() + sleep_seconds + 5 + while not orchestrator._top_models_ready.is_set() and time.monotonic() < deadline: + time.sleep(0.05) + + assert "unsloth/slow-GGUF" in orchestrator.default_models diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 2dd0d02515..f59a952b3a 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -81,11 +81,6 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { - exportConversationRawJsonl, - exportConversationCsv, - exportConversationShareGPT, -} from "@/features/chat/prompt-storage/prompt-storage-dialog"; import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, @@ -174,6 +169,36 @@ const TestTubeOutlineIcon = TestTube01Icon.slice( 3, ) as typeof TestTube01Icon; + +type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl"; + +const CHAT_EXPORT_OPTIONS: Array<{ + label: string; + format: ConversationExportFormat; +}> = [ + { label: "Raw JSONL", format: "raw-jsonl" }, + { label: "CSV", format: "csv" }, + { label: "ShareGPT JSONL", format: "sharegpt-jsonl" }, +]; + +async function exportConversationByFormat( + threadId: string, + format: ConversationExportFormat, +): Promise { + const exports = await import( + "@/features/chat/prompt-storage/prompt-storage-dialog" + ); + switch (format) { + case "raw-jsonl": + return exports.exportConversationRawJsonl(threadId); + case "csv": + return exports.exportConversationCsv(threadId); + case "sharegpt-jsonl": + return exports.exportConversationShareGPT(threadId); + } +} + + function runStatusDotClass(status: TrainingRunSummary["status"]): string { switch (status) { case "running": @@ -899,11 +924,7 @@ export function AppSidebar() { Export - {[ - { label: "Raw JSONL", fn: exportConversationRawJsonl }, - { label: "CSV", fn: exportConversationCsv }, - { label: "ShareGPT JSONL", fn: exportConversationShareGPT }, - ].map(({ label, fn }) => ( + {CHAT_EXPORT_OPTIONS.map(({ label, format }) => ( { @@ -911,7 +932,9 @@ export function AppSidebar() { const ids = item.type === "single" ? [item.id] : (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id); - await Promise.all(ids.map((id) => fn(id))); + await Promise.all( + ids.map((id) => exportConversationByFormat(id, format)), + ); } catch { toast.error("Export failed."); } diff --git a/studio/frontend/src/config/env.ts b/studio/frontend/src/config/env.ts index def1b9bad9..63cdd03141 100644 --- a/studio/frontend/src/config/env.ts +++ b/studio/frontend/src/config/env.ts @@ -54,6 +54,17 @@ export const usePlatformStore = create()((_, get) => ({ isChatOnly: () => get().chatOnly, })); +// Once an authoritative (server-reported) platform has been fetched, a +// non-forced response must not overwrite it. The post-render fetchDeviceType() +// in main.tsx runs before auth is ready and can resolve after the authed +// root-route/provider fetches; such a late write would reset deviceType, +// cloudflareUrl/serverUrl/secure, and fetched, whether it is a browser fallback +// (unauthenticated) or an earlier authenticated request that landed after a +// later forced refresh. Forced refreshes are explicit re-reads, so they still write. +function shouldKeepAuthoritativePlatform(force?: boolean): boolean { + return !force && usePlatformStore.getState().fetched; +} + // `force` re-reads /api/health even if cached, to pick up a late-arriving tunnel URL. export async function fetchDeviceType(options?: { force?: boolean; @@ -81,6 +92,15 @@ export async function fetchDeviceType(options?: { server_url?: string | null; secure?: boolean; }; + // Once the store holds an authoritative (server-reported) platform, a + // non-forced response must not overwrite it. It may be an unauthenticated + // fallback, or an earlier authenticated request that resolved after a + // later forced refresh already picked up device_type and the tunnel + // fields; writing either would reset device type or null the tunnel + // fields. Forced refreshes are explicit re-reads, so they still write. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = data.device_type ?? detectLocalPlatform(); const chatOnly = data.chat_only ?? false; const chatOnlyReason = data.chat_only_reason ?? null; @@ -101,7 +121,11 @@ export async function fetchDeviceType(options?: { } catch { // Backend not ready: use client-side detection so chat-only guard works // on initial load (important for macOS). Keep fetched=false so a later - // call retries against the backend. + // call retries against the backend. But a late non-forced failure must not + // wipe an authoritative platform that already resolved. + if (shouldKeepAuthoritativePlatform(options?.force)) { + return usePlatformStore.getState().deviceType; + } const deviceType = detectLocalPlatform(); const chatOnly = deviceType === "mac"; usePlatformStore.setState({ deviceType, chatOnly, fetched: false }); diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index cd7cfc77fc..b155eff780 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -35,7 +35,6 @@ import { useNativeModelDrop, useNativePathLeasesSupported, } from "@/features/native-intents"; -import { ProjectSourcesPanel } from "@/features/rag/components/project-sources-panel"; import { GuidedTour, useGuidedTourController } from "@/features/tour"; import { isTauri } from "@/lib/api-base"; import { toast } from "@/lib/toast"; @@ -51,7 +50,9 @@ import { Tooltip as TooltipPrimitive } from "radix-ui"; import { type CSSProperties, type ReactElement, + lazy, memo, + Suspense, useCallback, useEffect, useMemo, @@ -134,6 +135,13 @@ import { } from "./utils/chat-history-storage"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; + +const ProjectSourcesPanel = lazy(() => + import("@/features/rag/components/project-sources-panel").then((module) => ({ + default: module.ProjectSourcesPanel, + })), +); + type LoraCandidate = { id: string; baseModel: string; @@ -1018,7 +1026,15 @@ function ProjectLanding({ {projectTab === "sources" ? ( - + + Loading sources… + + } + > + + ) : (
{items.map((item) => { @@ -2246,12 +2262,29 @@ export function ChatPage({ return [...fromLoras, ...localModels]; }, [lorasFromStore, localModels]); - useEffect(() => { - if (getTrainingCompareHandoff()) return; - void refresh(); + const inventoryRefreshStartedRef = useRef(false); + const refreshDeferredModelInventories = useCallback(() => { + inventoryRefreshStartedRef.current = true; + void refresh({ includeLoras: true }); refreshLocalModels(); }, [refresh, refreshLocalModels]); + useEffect(() => { + if (getTrainingCompareHandoff()) return; + void refresh({ includeLoras: false }); + const timeoutId = window.setTimeout(() => { + if (!inventoryRefreshStartedRef.current) { + refreshDeferredModelInventories(); + } + }, 1200); + return () => window.clearTimeout(timeoutId); + }, [refresh, refreshDeferredModelInventories]); + + useEffect(() => { + if (!active || !modelSelectorOpen) return; + refreshDeferredModelInventories(); + }, [active, modelSelectorOpen, refreshDeferredModelInventories]); + useEffect(() => { // ChatPage no longer remounts on navigation, so re-check the handoff whenever // we return to /chat (e.g. from the training progress "compare in chat" action). diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index e23d1b0b33..798c1658f9 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -312,14 +312,18 @@ export function useChatModelRuntime() { [], ); - const refresh = useCallback(async (options?: { signal?: AbortSignal }) => { + const refresh = useCallback(async (options?: { + signal?: AbortSignal; + includeLoras?: boolean; + }) => { const signal = options?.signal; + const includeLoras = options?.includeLoras ?? true; setModelsError(null); try { const [listRes, statusRes, lorasRes] = await Promise.all([ listModels(), getInferenceStatus(), - listLoras(), + includeLoras ? listLoras() : Promise.resolve(null), ]); // Cancellation can land while the requests above are in flight. Bail @@ -327,7 +331,9 @@ export function useChatModelRuntime() { if (signal?.aborted) return; setModels(listRes.models.map(toChatModelSummary)); - setLoras(lorasRes.loras.map(toLoraSummary)); + if (lorasRes) { + setLoras(lorasRes.loras.map(toLoraSummary)); + } const selectedCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const isExternalSelectionActive = isExternalModelId(selectedCheckpoint); diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 67980b94c6..b545695f9e 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -22,7 +22,6 @@ import { unstable_useRemoteThreadListRuntime as useRemoteThreadListRuntime, } from "@assistant-ui/react"; import { createAssistantStream } from "assistant-stream"; -import mammoth from "mammoth"; import { type ReactElement, type ReactNode, @@ -33,7 +32,6 @@ import { useMemo, useRef, } from "react"; -import { extractText, getDocumentProxy } from "unpdf"; import { toast } from "sonner"; import { StudioWebSpeechDictationAdapter } from "./adapters/studio-web-speech-dictation-adapter"; import { @@ -181,7 +179,10 @@ class PDFAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const buffer = new Uint8Array(await attachment.file.arrayBuffer()); + const [{ extractText, getDocumentProxy }, buffer] = await Promise.all([ + import("unpdf"), + attachment.file.arrayBuffer().then((bytes) => new Uint8Array(bytes)), + ]); const pdf = await getDocumentProxy(buffer); const { text } = await extractText(pdf, { mergePages: true }); return { @@ -298,7 +299,10 @@ class DocxAttachmentAdapter implements AttachmentAdapter { } async send(attachment: PendingAttachment): Promise { - const arrayBuffer = await attachment.file.arrayBuffer(); + const [{ default: mammoth }, arrayBuffer] = await Promise.all([ + import("mammoth"), + attachment.file.arrayBuffer(), + ]); const { value } = await mammoth.extractRawText({ arrayBuffer }); return { id: attachment.id, diff --git a/studio/frontend/src/main.tsx b/studio/frontend/src/main.tsx index 0922e764bb..d0ddf2fc6e 100644 --- a/studio/frontend/src/main.tsx +++ b/studio/frontend/src/main.tsx @@ -5,8 +5,8 @@ import { StrictMode } from "react"; import { createRoot } from "react-dom/client"; import "./index.css"; -import { fetchDeviceType } from "./config/env"; import { App } from "./app/app"; +import { fetchDeviceType } from "./config/env"; import { initializeLocale } from "./i18n"; const globalCrypto = globalThis.crypto as Crypto | undefined; @@ -36,10 +36,10 @@ if (!rootElement) { initializeLocale(); -fetchDeviceType().then(() => { - createRoot(rootElement).render( - - - , - ); -}); +createRoot(rootElement).render( + + + , +); + +fetchDeviceType().catch(() => undefined); diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 5a48d26632..7ef5244754 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -498,19 +498,68 @@ mod tests { } #[cfg(unix)] - fn remove_managed_capability_cache() { - let _ = std::fs::remove_file( - dirs::home_dir() + static MANAGED_CAPABILITY_CACHE_TEST_LOCK: std::sync::LazyLock> = + std::sync::LazyLock::new(|| tokio::sync::Mutex::new(())); + + #[cfg(unix)] + struct ManagedCapabilityCacheHome { + path: PathBuf, + previous: Option, + } + + #[cfg(unix)] + impl ManagedCapabilityCacheHome { + fn new(test_name: &str) -> Self { + use std::time::{SystemTime, UNIX_EPOCH}; + + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) .unwrap() - .join(".unsloth") - .join("studio") - .join("desktop_capability_cache.json"), - ); + .as_nanos(); + let path = std::env::temp_dir().join(format!( + "unsloth-preflight-cache-{test_name}-{}-{nanos}", + std::process::id() + )); + std::fs::create_dir_all(&path).unwrap(); + let previous = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME"); + std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", &path); + Self { path, previous } + } + } + + #[cfg(unix)] + impl Drop for ManagedCapabilityCacheHome { + fn drop(&mut self) { + if let Some(previous) = &self.previous { + std::env::set_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME", previous); + } else { + std::env::remove_var("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME"); + } + let _ = std::fs::remove_dir_all(&self.path); + } + } + + #[cfg(unix)] + fn managed_capability_cache_path_for_test() -> PathBuf { + std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME") + .map(PathBuf::from) + .or_else(dirs::home_dir) + .unwrap() + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json") + } + + #[cfg(unix)] + fn remove_managed_capability_cache() { + let _ = std::fs::remove_file(managed_capability_cache_path_for_test()); } #[cfg(unix)] #[tokio::test] async fn managed_cli_capability_probe_classifies_core_cases() { + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("core-cases"); remove_managed_capability_cache(); for (name, script, stale_reason) in [ @@ -567,6 +616,75 @@ exit 1 } } + #[cfg(unix)] + #[tokio::test] + async fn managed_cli_capability_help_probe_runs_before_cache() { + use std::fs; + + let _cache_guard = MANAGED_CAPABILITY_CACHE_TEST_LOCK.lock().await; + let _cache_home = ManagedCapabilityCacheHome::new("cache-hit"); + + remove_managed_capability_cache(); + // `-h` always succeeds unless `modeh` exists; the desktop-capabilities + // probe always succeeds unless `modecap` exists. Toggling those lets us + // prove the ordering: -h runs on every probe (even a cache hit), while + // the heavier capability probe is skipped once the cache is warm. + let fake = fake_cli( + "cap-cache-hit", + r#"#!/bin/sh +log="$0.calls" +modeh="$0.modeh" +modecap="$0.modecap" +printf '%s\n' "$*" >> "$log" +if [ "$1" = "-h" ]; then + if [ -f "$modeh" ]; then exit 42; fi + exit 0 +fi +if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then + if [ -f "$modecap" ]; then exit 42; fi + printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}' + exit 0 +fi +exit 1 +"#, + ); + let bin = fake.bin.clone(); + let calls = bin.with_extension("calls"); + let modeh = bin.with_extension("modeh"); + let modecap = bin.with_extension("modecap"); + + // Cold probe: runs -h and the capability probe, then caches the result. + assert!(matches!( + probe_managed_bin(bin.clone()).await, + ManagedProbe::Ready { .. } + )); + let first_calls = fs::read_to_string(&calls).unwrap(); + assert!(first_calls.contains("-h")); + assert!(first_calls.contains("studio desktop-capabilities --json")); + + // Cache hit: -h still runs, but the capability probe is skipped (breaking + // it via `modecap` proves it is not invoked). + fs::write(&modecap, "broken").unwrap(); + fs::write(&calls, "").unwrap(); + assert!(matches!( + probe_managed_bin(bin.clone()).await, + ManagedProbe::Ready { .. } + )); + assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + + // A non-launchable CLI is caught by the -h probe even with a warm cache: + // preflight reports Stale (for repair) and never trusts the cache. + fs::write(&modeh, "broken").unwrap(); + fs::write(&calls, "").unwrap(); + assert!(matches!( + probe_managed_bin(bin).await, + ManagedProbe::Stale { .. } + )); + assert_eq!(fs::read_to_string(&calls).unwrap(), "-h\n"); + + remove_managed_capability_cache(); + } + const EXPECTED_ROOT_ID: &str = "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"; const OTHER_ROOT_ID: &str = "bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb"; diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 57b8365ec5..0d20f271c5 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -188,6 +188,16 @@ fn managed_bin_fingerprint(bin: &Path) -> Option { } fn capability_cache_path() -> Option { + #[cfg(test)] + if let Some(home) = std::env::var_os("UNSLOTH_TEST_DESKTOP_CAPABILITY_CACHE_HOME") { + return Some( + PathBuf::from(home) + .join(".unsloth") + .join("studio") + .join("desktop_capability_cache.json"), + ); + } + dirs::home_dir().map(|home| { home.join(".unsloth") .join("studio") @@ -400,6 +410,12 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool { pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe { let started = Instant::now(); + // Always verify the managed CLI actually launches before trusting the cache. + // A matching capability fingerprint does not prove the binary can still run: + // its venv interpreter or a runtime dependency can be broken while the + // path/size/mtime/markers are unchanged, so the -h probe runs first and a + // non-launchable install is reported Stale for repair. The capability cache + // below still skips the heavier desktop-capabilities probe on a hit. if !run_cli_probe(&bin, &["-h"]).await { info!( "Managed preflight: cli unusable for {:?} in {}ms", From e7e6a0fb475963e9747358e84cb0a22994c1a54a Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 01:50:35 -0700 Subject: [PATCH 008/402] Polish assistant message actions menu (#6962) * Polish assistant message actions menu Use the circle question mark (HelpCircleIcon) for the "See response details" action instead of the file-database icon, and lowercase the "Export as markdown" label. * Align response details sheet icon --- .../assistant-ui/message-response-details-sheet.tsx | 4 ++-- studio/frontend/src/components/assistant-ui/thread.tsx | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx index 823696693a..331a06a4c6 100644 --- a/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx +++ b/studio/frontend/src/components/assistant-ui/message-response-details-sheet.tsx @@ -18,7 +18,7 @@ import { useExternalProvidersStore, } from "@/features/chat"; import { cn } from "@/lib/utils"; -import { FileDatabaseIcon } from "@hugeicons/core-free-icons"; +import { HelpCircleIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { useMessage, useMessageTiming } from "@assistant-ui/react"; import type { FC, ReactNode } from "react"; @@ -341,7 +341,7 @@ export const MessageResponseDetailsSheet: FC<{ diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 09551cd413..d987092c48 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -127,6 +127,7 @@ import { FileDatabaseIcon, Folder01Icon, FolderAddIcon, + HelpCircleIcon, Image03Icon, McpServerIcon, PencilRulerIcon, @@ -3952,7 +3953,7 @@ const AssistantActionBar: FC = () => { strokeWidth={1.75} className="size-icon" /> - Export as Markdown + Export as markdown { className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground" > From 7f9964f21ed540fdf1ecc1549947bfa5353e6127 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 01:51:34 -0700 Subject: [PATCH 009/402] Move New badge to System settings tab (#6963) * Move New badge to System settings tab Show the "New" badge on the System tab and drop it from Connections. * Stabilize refresh revocation UI test * [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> --- .../src/features/settings/settings-dialog.tsx | 2 +- tests/studio/playwright_chat_ui.py | 58 ++++++++++++++++--- 2 files changed, 50 insertions(+), 10 deletions(-) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 0d201edc0d..d8a0d45d3f 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -56,6 +56,7 @@ const TABS: TabDef[] = [ id: "resources", labelKey: "settings.tabs.resources", icon: CpuIcon, + badgeKey: "common.new", }, { id: "chat", @@ -72,7 +73,6 @@ const TABS: TabDef[] = [ id: "connections", labelKey: "settings.tabs.connections", icon: CloudIcon, - badgeKey: "common.new", }, { id: "about", labelKey: "settings.tabs.about", icon: HelpCircleIcon }, ]; diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index a53534acc0..71297f9043 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -380,6 +380,18 @@ with sync_playwright() as p: fail(f"/api/auth/refresh wedged: {refresh_resp['error']!r}") refresh = refresh_resp.get("body") or {} token = (refresh or {}).get("access_token") + next_refresh_token = (refresh or {}).get("refresh_token") + if token and next_refresh_token: + robust_evaluate( + page, + """([accessToken, refreshToken]) => { + localStorage.setItem('unsloth_auth_token', accessToken); + localStorage.setItem('unsloth_auth_refresh_token', refreshToken); + }""", + [token, next_refresh_token], + ) + elif token: + fail("/api/auth/refresh returned access_token but no refresh_token") if not token: fail("could not obtain auth token after change-password") @@ -1169,6 +1181,13 @@ with sync_playwright() as p: fail(f"curl login returned no access_token: {login_body!r}") info("CLI obtained an access token") + browser_refresh_token = robust_evaluate( + page, + "() => localStorage.getItem('unsloth_auth_refresh_token')", + ) + if not browser_refresh_token: + fail("browser refresh token missing before CLI rotation") + change_proc = subprocess.run( [ "curl", @@ -1203,18 +1222,39 @@ with sync_playwright() as p: # /change-password revoked refresh tokens server-side (auth.py), so # the browser's /api/auth/refresh must now fail. - refresh_after = evaluate_fetch( - page, - f"{BASE}/api/auth/refresh", - method = "POST", - timeout_ms = FETCH_TIMEOUT_MS, + refresh_proc = subprocess.run( + [ + "curl", + "-sS", + "-o", + os.devnull, + "-w", + "%{http_code}", + "-X", + "POST", + f"{BASE}/api/auth/refresh", + "-H", + "Content-Type: application/json", + "-d", + json.dumps({"refresh_token": browser_refresh_token}), + ], + capture_output = True, + text = True, + timeout = 15, ) - if refresh_after.get("error"): - fail(f"/api/auth/refresh wedged: {refresh_after['error']!r}") - if refresh_after["status"] == 200: + if refresh_proc.returncode != 0: + fail( + f"curl refresh-token check failed: rc={refresh_proc.returncode} " + f"stderr={refresh_proc.stderr!r} stdout={refresh_proc.stdout!r}" + ) + try: + refresh_status = int(refresh_proc.stdout.strip()) + except ValueError: + fail(f"curl refresh-token check returned invalid status: " f"{refresh_proc.stdout!r}") + if refresh_status == 200: fail(f"/api/auth/refresh should fail after CLI rotation; got 200") info( - f"OK browser /api/auth/refresh now {refresh_after['status']} " + f"OK browser /api/auth/refresh now {refresh_status} " "(refresh token revoked) -- old studio session can no longer renew" ) From 393d7e9c2b48b928126008399b0c9443af8a9ec7 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:22:36 +0100 Subject: [PATCH 010/402] Fix opencode Unsloth provider selection (#6906) * fix: force Unsloth provider selection for opencode * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * opencode: pin the model without clobbering the user's disabled providers The session overlay wrote disabled_providers unconditionally and the inline OPENCODE_CONFIG_CONTENT set disabled_providers to an empty list. Since that inline layer outranks the user's global and project config and opencode replaces the array rather than merging it, every provider the user had disabled was silently re-enabled for the session. Only strip 'unsloth' from an existing disable list, and drop disabled_providers from the inline config. Also insert --model only on a bare launch: it is a global flag for the TUI, so placing it before a passthrough subcommand (serve/run) breaks arg parsing; a subcommand takes the model from the pinned config instead. Parse the printed OPENCODE_CONFIG_CONTENT with shlex.split in the test so it round-trips under POSIX shell quoting. * Re-enable a globally disabled opencode unsloth provider for the session A fresh OPENCODE_CONFIG overlay omits disabled_providers, and opencode replaces that array across config layers only when a higher layer sets the key, so a user's global disabled_providers of ['unsloth', ...] survived the merge and left the session provider disabled even though the overlay defines provider.unsloth and pins the model. Consult the user's global opencode config (XDG_CONFIG_HOME/opencode, or %APPDATA%/opencode on Windows) when the overlay has no list of its own, and when the effective list disables unsloth write it back to the overlay minus unsloth. The provider loads while the user's other disabled providers stay disabled. Best-effort read: a missing or unparseable global config is a no-op. * Override opencode disabled_providers in the inline layer; keep model flag for TUI flags Re-enabling a disabled unsloth provider now rides in the inline OPENCODE_CONFIG_CONTENT layer instead of the session overlay. The overlay sits below a project opencode.json, which could re-disable the provider; the inline layer outranks both global and project configs and is recomputed each run, so no-launch reruns never reuse a stale generated list. The effective disabled list is read from the project config if the repo sets one, else the global config, across config.json/opencode.json/opencode.jsonc (JSONC tolerated), and written back minus unsloth only when unsloth is disabled. Also keep the pinned --model when the opencode passthrough starts with a top-level TUI flag such as --dir or --continue; only a real subcommand (serve/run/...) takes the model from config, so a leading '-' now still gets --model injected. * Discover the opencode project config by walking up from the cwd opencode finds a project config by searching ancestor directories, not just the cwd. Walk from the cwd up to the filesystem root and use the nearest directory that sets disabled_providers, so running unsloth start opencode from a subdirectory of a repo whose root config disables unsloth still gets the inline override. * Only inject opencode --model on a bare launch; rely on the inline model pin Injecting --model whenever the passthrough started with a flag could place it before a subcommand (e.g. opencode --print-logs serve), which opencode can misparse. --model is unnecessary for any passthrough because the inline OPENCODE_CONFIG_CONTENT pins the model in the highest-priority layer, so the session model is forced without the flag. Restrict --model to the bare launch and pass any other invocation through untouched. * Register the session provider under a dedicated OpenCode id Selecting the Unsloth model reliably required the wrapper to re-enable a user-disabled unsloth provider, which meant reconstructing OpenCode's full disabled_providers resolution (global, OPENCODE_CONFIG overlay, project config discovered via --dir or an ancestor walk, .opencode directories, OPENCODE_CONFIG_DIR, config.json/opencode.json/opencode.jsonc precedence, and {env:} variable substitution) and overriding it in the inline layer. That is unbounded and cannot be kept correct. Register the session provider under a dedicated id (unsloth-studio) instead. A user's disabled_providers list would never target it, so the session model is always selectable and the overlay no longer reads or writes disabled_providers at all: the user's own disables, in whatever config layer, are left exactly as they are. This removes the JSONC parser, the config-directory scan, and the ancestor/global resolution helpers, and the tests that exercised them. * Scope the opencode session to the Studio provider opencode filters every provider, including a config-defined custom one, through its enabled_providers allowlist and disabled_providers denylist, and pinning the model does not bypass that gate (a filtered provider resolves to a not-found error). The provider arrays are also replaced, not merged, across config layers. So a user with an enabled_providers allowlist that omits the session provider would still have the Studio model filtered out. Set enabled_providers to just the session provider and clear disabled_providers in the inline OPENCODE_CONFIG_CONTENT overlay (the highest-priority layer, which replaces these arrays). This guarantees the Studio model loads regardless of the user's provider filters, without reading or reconstructing their multi-layer config. It is session-only: the overlay lives in the env for this launch and never touches the user's config files, so their normal opencode is unchanged and only this session is limited to the Studio provider. Also drop the redundant --model on --no-launch so the printed command stays append-safe for drivers that append a subcommand (the inline pin forces the model), and parse both POSIX and PowerShell no-launch output in the opencode tests so they are not shell-specific. * Pin opencode small_model to the session provider The session allowlists only the Studio provider, but opencode's separate small_model (used for lightweight tasks) could still point at another provider from the user or project config; under the allowlist that provider is filtered, so the lightweight task would resolve a not-found error mid-session even with the main model pinned. Pin small_model to the session model in the same inline overlay so every model use stays on the enabled provider. The session serves one model, so it is the only valid target, and this stays session-only like the rest of the overlay. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Wasim Yousef Said --- unsloth_cli/commands/start.py | 54 +++++++++++-- unsloth_cli/tests/test_start.py | 132 ++++++++++++++++++++++++++------ 2 files changed, 158 insertions(+), 28 deletions(-) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 1895125b11..d959d20c83 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -50,6 +50,12 @@ _HERMES_PROVIDER = "unsloth" # windows and scales the compaction threshold back down to the real window. _HERMES_MIN_CONTEXT = 65536 _PI_PROVIDER = "unsloth" +# OpenCode selects a model by "/" and honors a user +# disabled_providers list. Register the session provider under a dedicated id a +# user's disable list would never target, so the model is always selectable +# without the wrapper having to reconstruct (and override) OpenCode's full, +# multi-layer disabled_providers resolution. +_OPENCODE_PROVIDER = "unsloth-studio" _PROVIDER_HEADER = f"[model_providers.{_CODEX_PROFILE}]" _PASSTHROUGH = {"allow_extra_args": True, "ignore_unknown_options": True} _CLAUDE_ENV_UNSET = ("ANTHROPIC_API_KEY", "CLAUDE_CODE_OAUTH_TOKEN") @@ -1116,13 +1122,17 @@ def write_opencode_config( config = _read_json_object(path) if config is None: typer.echo( - f"Warning: couldn't parse {path} — add an 'unsloth' provider there " - "yourself, or move the file aside and re-run.", + f"Warning: couldn't parse {path} — add an '{_OPENCODE_PROVIDER}' provider " + "there yourself, or move the file aside and re-run.", err = True, ) return {} before = json.dumps(config, sort_keys = True) config.setdefault("$schema", "https://opencode.ai/config.json") + # The session provider is registered under a dedicated id (_OPENCODE_PROVIDER) + # that a user's disabled_providers list would never target, so it is always + # selectable without this overlay having to reconstruct or override OpenCode's + # disabled_providers resolution. model_entry = {"name": model["id"]} window = model.get("context_length") or model.get("max_context_length") if window: @@ -1131,14 +1141,14 @@ def write_opencode_config( # disables OpenCode's auto-compaction; declare the real window (and a sane # output cap) so it compacts instead of overflowing the server. model_entry["limit"] = {"context": window, "output": min(window // 4, 8192)} - _subdict(config, "provider")["unsloth"] = { + _subdict(config, "provider")[_OPENCODE_PROVIDER] = { "npm": "@ai-sdk/openai-compatible", "name": "Unsloth Studio", "options": {"baseURL": f"{base}/v1", "apiKey": key}, "models": {model["id"]: model_entry}, } # OpenCode selects a model by "/". - config["model"] = f"unsloth/{model['id']}" + config["model"] = f"{_OPENCODE_PROVIDER}/{model['id']}" if window: # Compact with ~10% headroom (near 90% full). The fixed 20k-token default # buffer over-compacts, or never settles, on a small local context. @@ -1450,7 +1460,20 @@ def opencode( serve = serve, launch = launch, ) - command = ["opencode", *ctx.args] + opencode_model = f"{_OPENCODE_PROVIDER}/{entry['id']}" + # The inline OPENCODE_CONFIG_CONTENT below pins the model in the highest-priority + # layer, so the session model is forced without a --model flag. Only add --model for + # an interactive bare launch (a convenience so the TUI opens on our model). It is + # omitted for passthrough (inserting it before a subcommand can be misparsed) and for + # --no-launch, where the printed command is consumed by drivers that append a + # subcommand such as `run `; a leading --model would land before that + # subcommand and break it. Those paths rely on the inline pin instead. + if ctx.args: + command = ["opencode", *ctx.args] + elif launch: + command = ["opencode", "--model", opencode_model] + else: + command = ["opencode"] with _session_config("opencode", launch) as cfg: config_path = cfg / "opencode.json" # OPENCODE_CONFIG is an overlay (loaded between the user's global and project @@ -1462,7 +1485,26 @@ def opencode( # outranks project config; the API key stays in the private file, never the env. # Only --yolo carries a permission here (its allow must win over a project config); # a non-yolo session returns no permission, so the project's own rules are honored. - inline_config: dict = {"model": f"unsloth/{entry['id']}"} + # opencode filters every provider (a config-defined custom one included) through + # its enabled_providers allowlist and disabled_providers denylist, and a model pin + # does not bypass that gate -- a filtered provider resolves to ModelNotFoundError. + # To guarantee the session model loads without reading or modifying the user's real + # config, scope THIS session to our provider alone: allowlist _OPENCODE_PROVIDER and + # clear the denylist. These arrays are replaced (not merged) by higher layers, so + # setting them in the highest-priority inline overlay neutralizes any user allowlist + # or denylist for the launch. It is session-only: it lives in OPENCODE_CONFIG_CONTENT + # for this invocation and never touches the user's config files, so their normal + # `opencode` is unchanged; only this session is limited to the Studio provider. + # small_model is opencode's separate model for lightweight tasks; pin it to the + # session model too, or a user/project small_model on another (now filtered) + # provider would resolve a not-found error mid-session. The session serves one + # model, so the session model is the only valid target here anyway. + inline_config: dict = { + "model": opencode_model, + "small_model": opencode_model, + "enabled_providers": [_OPENCODE_PROVIDER], + "disabled_providers": [], + } if session_permission: inline_config["permission"] = session_permission env = { diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index affc24626e..58888010b3 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -435,15 +435,10 @@ def test_opencode_inline_config_beats_project_config(fake_studio): # permissions) ride in OPENCODE_CONFIG_CONTENT, which outranks project config. result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--yolo"]) assert result.exit_code == 0, result.output - content_line = next( - ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") - ) - inline = json.loads( - shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] - ) - assert inline["model"] == f"unsloth/{MODEL['id']}" + inline = _opencode_inline_config(result.output) + assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" assert inline["permission"] == {"edit": "allow", "bash": "allow", "webfetch": "allow"} - assert "sk-unsloth" not in content_line # key stays in the private file + assert "sk-unsloth" not in result.output # key stays in the private file, not the env def test_opencode_inline_config_omits_permission_without_yolo(fake_studio): @@ -452,13 +447,8 @@ def test_opencode_inline_config_omits_permission_without_yolo(fake_studio): # user's project rules; clearing our own config is the fix, and the inline pins the model. result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output - content_line = next( - ln for ln in result.output.splitlines() if ln.startswith("export OPENCODE_CONFIG_CONTENT=") - ) - inline = json.loads( - shlex.split(content_line.removeprefix("export OPENCODE_CONFIG_CONTENT="))[0] - ) - assert inline["model"] == f"unsloth/{MODEL['id']}" + inline = _opencode_inline_config(result.output) + assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" assert "permission" not in inline @@ -1376,14 +1366,17 @@ def test_write_opencode_config_fresh(tmp_path): path = tmp_path / "opencode.json" start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) config = json.loads(path.read_text()) - provider = config["provider"]["unsloth"] + provider = config["provider"][start._OPENCODE_PROVIDER] assert provider["npm"] == "@ai-sdk/openai-compatible" assert provider["options"] == {"baseURL": f"{BASE}/v1", "apiKey": "sk-unsloth-abc"} # Context limit must be declared, or OpenCode treats it as 0 and disables compaction. assert provider["models"] == { MODEL["id"]: {"name": MODEL["id"], "limit": {"context": 131072, "output": 8192}} } - assert config["model"] == f"unsloth/{MODEL['id']}" + assert config["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + # The overlay never writes disabled_providers; the dedicated provider id is one a + # user's disable list would not target, so nothing needs re-enabling. + assert "disabled_providers" not in config # Compaction buffer scaled to ~10% of the window (compact near 90%). assert config["compaction"] == {"auto": True, "reserved": 131072 // 10} @@ -1391,18 +1384,98 @@ def test_write_opencode_config_fresh(tmp_path): def test_write_opencode_config_preserves_and_idempotent(tmp_path): path = tmp_path / "opencode.json" path.write_text( - json.dumps({"theme": "tokyonight", "provider": {"anthropic": {"name": "Anthropic"}}}) + json.dumps( + { + "theme": "tokyonight", + "disabled_providers": ["ollama", "unsloth"], + "provider": {"anthropic": {"name": "Anthropic"}}, + } + ) ) start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) config = json.loads(path.read_text()) assert config["theme"] == "tokyonight" + # The overlay no longer edits disabled_providers; re-enabling unsloth is done in + # the inline layer, so an existing list here is preserved untouched. + assert config["disabled_providers"] == ["ollama", "unsloth"] assert config["provider"]["anthropic"]["name"] == "Anthropic" - assert config["provider"]["unsloth"]["options"]["baseURL"] == f"{BASE}/v1" + assert config["provider"][start._OPENCODE_PROVIDER]["options"]["baseURL"] == f"{BASE}/v1" before = path.read_text() start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) assert path.read_text() == before +def test_write_opencode_config_keeps_foreign_disabled_providers(tmp_path): + # A user who disabled other providers (but not unsloth) must keep them disabled: + # the overlay must not rewrite disabled_providers, or those providers get silently + # re-enabled for the session. + path = tmp_path / "opencode.json" + path.write_text(json.dumps({"disabled_providers": ["openai", "gemini"]})) + start.write_opencode_config(BASE, "sk-unsloth-abc", MODEL, path) + config = json.loads(path.read_text()) + assert config["disabled_providers"] == ["openai", "gemini"] + + +def _opencode_inline_config(output: str) -> dict: + # --no-launch prints OPENCODE_CONFIG_CONTENT as a POSIX `export NAME=` + # line on Unix/WSL and a PowerShell `$env:NAME = ""` line on native Windows; + # parse whichever the host emitted so the opencode tests are shell-agnostic. + name = "OPENCODE_CONFIG_CONTENT" + for raw in output.splitlines(): + line = raw.strip() + if line.startswith(f"export {name}="): + return json.loads(shlex.split(line.removeprefix(f"export {name}="))[0]) + prefix = f'$env:{name} = "' + if line.startswith(prefix) and line.endswith('"'): + escaped = line[len(prefix) : -1] + # Reverse _print_env's PowerShell escaping (backtick is the escape char). + value = escaped.replace("`$", "$").replace('`"', '"').replace("``", "`") + return json.loads(value) + raise AssertionError(f"{name} not found in:\n{output}") + + +def test_opencode_inline_scopes_session_to_studio_provider(fake_studio): + # opencode filters even config-defined providers through enabled/disabled_providers, + # and a model pin does not bypass that gate. The inline overlay (session-only, highest + # layer, arrays replace) allowlists our provider and clears the denylist so the Studio + # model always loads regardless of the user's config, without reading or editing it. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) + assert result.exit_code == 0, result.output + inline = _opencode_inline_config(result.output) + assert inline["enabled_providers"] == [start._OPENCODE_PROVIDER] + assert inline["disabled_providers"] == [] + assert inline["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + # small_model stays on the enabled provider too, so lightweight tasks do not resolve a + # filtered provider mid-session. + assert inline["small_model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + + +def test_opencode_passthrough_flags_omit_model_flag(fake_studio): + # Any passthrough (top-level flags that may precede a subcommand, or a subcommand) + # is left untouched; --model is not injected. The model is pinned by the inline + # OPENCODE_CONFIG_CONTENT (highest layer) instead, so it is still forced. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "--dir", "repo"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command == ["opencode", "--dir", "repo"] + assert "--model" not in command + assert ( + _opencode_inline_config(result.output)["model"] + == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + ) + + +def test_opencode_passthrough_subcommand_omits_model_flag(fake_studio): + # A passthrough subcommand (e.g. `serve`) takes the model from the pinned config; + # inserting --model before it would break opencode's arg parsing. + result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch", "serve"]) + assert result.exit_code == 0, result.output + command = _launch_command(result.output) + assert command[0] == "opencode" + assert command[1] == "serve" + assert "--model" not in command + + def test_connect_opencode_no_launch(fake_studio, tmp_path): result = CliRunner().invoke(start.start_app, ["opencode", "--no-launch"]) assert result.exit_code == 0, result.output @@ -1410,9 +1483,24 @@ def test_connect_opencode_no_launch(fake_studio, tmp_path): config_path = tmp_path / "agents" / "opencode" / "opencode.json" # OPENCODE_CONFIG overlay points at the session file, not the user's global config. _assert_env_set(result.output, "OPENCODE_CONFIG", str(config_path)) + inline_config = _opencode_inline_config(result.output) config = json.loads(config_path.read_text()) - assert config["provider"]["unsloth"]["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" - assert config["model"] == f"unsloth/{MODEL['id']}" + provider = config["provider"][start._OPENCODE_PROVIDER] + assert provider["options"]["apiKey"] == "sk-unsloth-feedfacefeedface" + assert config["model"] == f"{start._OPENCODE_PROVIDER}/{MODEL['id']}" + # The session config file (a throwaway overlay, not the user's real config) does not + # carry provider filters; the session scoping rides in the inline env layer only. + assert "disabled_providers" not in config + assert "enabled_providers" not in config + assert inline_config == { + "model": f"{start._OPENCODE_PROVIDER}/{MODEL['id']}", + "small_model": f"{start._OPENCODE_PROVIDER}/{MODEL['id']}", + "enabled_providers": [start._OPENCODE_PROVIDER], + "disabled_providers": [], + } + # --no-launch prints an append-safe base command (no --model before a subcommand a + # driver may append); the model is forced by the inline pin above. + assert _launch_command(result.output) == ["opencode"] assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) @@ -1765,7 +1853,7 @@ def test_no_launch_rerun_clears_stale_opencode_yolo_permissions(fake_studio, tmp # revert to OpenCode's permissive "allow" default). assert config["permission"] == {"edit": "ask", "bash": "ask", "webfetch": "ask"} # The session provider survives the cleanup. - assert "unsloth" in config["provider"] + assert start._OPENCODE_PROVIDER in config["provider"] def test_no_launch_rerun_clears_stale_openclaw_yolo_state(fake_studio, tmp_path): From baacbd025d9ce92034cd3a41ad0d1ceacabf8202 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:23:25 +0100 Subject: [PATCH 011/402] Fix Hermes install hint on Windows (#6903) * fix: use Windows Hermes installer from unsloth start * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Skip the Hermes setup wizard during unattended start-install unsloth start hermes auto-installs Hermes and then writes its own session-scoped Hermes config. The install commands, as written, drop into the installer's interactive setup wizard (hermes setup), which prompts for global API keys and model choice and points the user at a different global provider than the one Unsloth just configured, blocking the launch. Pass the installer's skip flag on both platforms: the PowerShell scriptblock form with -SkipSetup, and bash -s -- --skip-setup for the piped POSIX installer. * Refresh PATH from the registry after a Windows agent install A Windows installer persists the agent's directory to the User/Machine PATH in the registry and updates only its own process, so the current process keeps a stale PATH until it restarts (the installers print 'restart your terminal'). The post-install shutil.which then misses the just-installed agent and unsloth start fails with 'installed but isn't on PATH yet', forcing a re-run in a new shell. Merge the registry PATH hives back into the process before re-resolving so a freshly installed agent launches in the same invocation. No-op off Windows and on any read error; only ever augments PATH. * Fix/adjust PATH refresh for PR #6903 --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- unsloth_cli/commands/start.py | 73 +++++++++++++++++++++++-- unsloth_cli/tests/test_start.py | 96 +++++++++++++++++++++++++++++++++ 2 files changed, 165 insertions(+), 4 deletions(-) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index d959d20c83..01014de4d3 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -44,6 +44,19 @@ _CODEX_PROFILE = "unsloth_api" _CODEX_ENV_KEY = "UNSLOTH_STUDIO_AUTH_TOKEN" _HERMES_ENV_KEY = "UNSLOTH_API_KEY" _HERMES_PROVIDER = "unsloth" +# Skip the installer's interactive setup wizard: `unsloth start hermes` runs +# this hint unattended and then writes its own session-scoped Hermes config, so +# the wizard's global API-key/model prompts would block the launch and point the +# user at a different (global) provider than the one Unsloth just configured. +# Both installers expose a skip flag: `-SkipSetup` (PowerShell) and +# `--skip-setup` (POSIX; passed to the piped script via `bash -s --`). +_HERMES_WINDOWS_INSTALL_HINT = ( + "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup" +) +_HERMES_POSIX_INSTALL_HINT = ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash -s -- --skip-setup" +) # Hermes refuses to initialize when the model window is under 64,000 tokens; its # error message points at the model.context_length / auxiliary.compression # overrides in config.yaml. write_hermes_config claims this value for smaller @@ -138,6 +151,10 @@ def _yolo_command_flags(agent: str, yolo: bool) -> list: return _YOLO_COMMAND_FLAGS.get(agent, []) if yolo else [] +def _hermes_install_hint() -> str: + return _HERMES_WINDOWS_INSTALL_HINT if os.name == "nt" else _HERMES_POSIX_INSTALL_HINT + + class LoadOptions(NamedTuple): """Model-load knobs forwarded to /api/inference/load when --model triggers a load.""" @@ -848,6 +865,54 @@ def _print_env( typer.echo(" ".join((*inline, shlex.join(command)))) +def _refresh_windows_path() -> None: + # Merge Windows registry PATH hives after the current process PATH so a + # freshly installed agent is visible without changing existing precedence. + if os.name != "nt": + return + try: + import winreg + except Exception: + return + + entries = [] + seen = set() + + def add_path(value: str) -> bool: + added = False + for entry in str(value).split(os.pathsep): + entry = entry.strip() + if not entry: + continue + key = os.path.normcase(entry).casefold() + if key in seen: + continue + seen.add(key) + entries.append(entry) + added = True + return added + + add_path(os.environ.get("PATH", "")) + added_registry = False + hives = ( + (winreg.HKEY_CURRENT_USER, "Environment"), + ( + winreg.HKEY_LOCAL_MACHINE, + r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", + ), + ) + for root, sub in hives: + try: + with winreg.OpenKey(root, sub) as key: + value, _ = winreg.QueryValueEx(key, "Path") + except OSError: + continue + if value: + added_registry = add_path(os.path.expandvars(str(value))) or added_registry + if added_registry: + os.environ["PATH"] = os.pathsep.join(entries) + + def _install_agent(name: str, install_hint: str) -> Optional[str]: # Missing agent under --launch: offer to run its documented install command, then # re-resolve it on PATH. Consent-based (we never auto-run a remote install script @@ -866,6 +931,9 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]: install_command = ["/bin/sh", "-c", install_hint] if subprocess.run(install_command).returncode != 0: _fail(f"Install command failed. Run it yourself, then re-run: {install_hint}") + # The installer just wrote PATH to the registry (Windows); pull it into this + # process so the freshly installed agent resolves without a shell restart. + _refresh_windows_path() executable = shutil.which(name) if executable is None: _fail( @@ -1536,10 +1604,7 @@ def hermes( launch = launch, ) command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] - install_hint = ( - "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" - "/main/scripts/install.sh | bash" - ) + install_hint = _hermes_install_hint() with _session_config("hermes", launch) as home: # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 58888010b3..bd964d5e54 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -92,6 +92,7 @@ def test_claude_flags_detected_when_version_not_first_token(monkeypatch): def test_install_agent_prompts_then_installs(monkeypatch): # TTY + yes: run the documented install command, then re-resolve the now-present binary. + monkeypatch.setattr(start.os, "name", "posix") monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) ran = [] @@ -108,6 +109,101 @@ def test_install_agent_prompts_then_installs(monkeypatch): assert ran == [["/bin/sh", "-c", "npm install -g @openai/codex"]] +def test_install_agent_uses_powershell_on_windows(monkeypatch): + monkeypatch.setattr(start.os, "name", "nt") + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: True) + ran = [] + monkeypatch.setattr( + start.subprocess, + "run", + lambda command, *a, **k: ran.append(command) or SimpleNamespace(returncode = 0), + ) + monkeypatch.setattr(start.shutil, "which", lambda _: r"C:\Users\samle\bin\hermes.exe") + + install_hint = "& ([scriptblock]::Create((irm https://x/install.ps1))) -SkipSetup" + executable = start._install_agent("hermes", install_hint) + + assert executable == r"C:\Users\samle\bin\hermes.exe" + assert ran == [["powershell", "-NoProfile", "-Command", install_hint]] + + +def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch): + monkeypatch.setattr(start.os, "name", "nt") + + # Scriptblock form so `-SkipSetup` reaches the installer and the interactive + # setup wizard is skipped during the unattended `unsloth start hermes` run. + assert start._hermes_install_hint() == ( + "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1)))" + " -SkipSetup" + ) + + +def test_hermes_install_hint_is_bash_on_posix(monkeypatch): + monkeypatch.setattr(start.os, "name", "posix") + + # `bash -s -- --skip-setup` forwards the skip flag to the piped installer. + assert start._hermes_install_hint() == ( + "curl -fsSL https://raw.githubusercontent.com/NousResearch/hermes-agent" + "/main/scripts/install.sh | bash -s -- --skip-setup" + ) + + +def test_refresh_windows_path_noop_off_windows(monkeypatch): + monkeypatch.setattr(start.os, "name", "posix") + before = os.environ.get("PATH", "") + monkeypatch.setenv("PATH", before) + start._refresh_windows_path() + assert os.environ.get("PATH", "") == before + + +def test_refresh_windows_path_merges_registry_hives(monkeypatch): + # Fake Windows registry PATH values written after this process started. + hkcu, hklm = object(), object() + reg = { + (hkcu, "Environment"): r"C:\existing;C:\Users\me\hermes\bin", + ( + hklm, + r"SYSTEM\CurrentControlSet\Control\Session Manager\Environment", + ): r"C:\Windows\System32", + } + + class _Key: + def __init__(self, value): + self._value = value + + def __enter__(self): + return self + + def __exit__(self, *a): + return False + + def open_key(root, sub): + if (root, sub) in reg: + return _Key(reg[(root, sub)]) + raise OSError("missing hive") + + fake_winreg = SimpleNamespace( + HKEY_CURRENT_USER = hkcu, + HKEY_LOCAL_MACHINE = hklm, + OpenKey = open_key, + QueryValueEx = lambda key, name: (key._value, 1), + ) + monkeypatch.setattr(start.os, "name", "nt") + monkeypatch.setattr(start.os, "pathsep", ";") + monkeypatch.setitem(sys.modules, "winreg", fake_winreg) + monkeypatch.setenv("PATH", r"C:\custom;C:\existing") + + start._refresh_windows_path() + + assert os.environ["PATH"].split(";") == [ + r"C:\custom", + r"C:\existing", + r"C:\Users\me\hermes\bin", + r"C:\Windows\System32", + ] + + def test_install_agent_declined_returns_none(monkeypatch): # TTY + no: never runs anything; caller falls back to the print-hint failure. monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) From a113f893ea767a701b589cd321408f718389a2b8 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Wed, 8 Jul 2026 06:30:37 -0300 Subject: [PATCH 012/402] Studio: heal DiffusionGemma tool calls into structured tool_calls (#6851) * Studio: heal DiffusionGemma tool calls into structured tool_calls * Fall back to supports_tools for backends without the passthrough capability * Route DiffusionGemma client tools through passthrough when enable_tools is on * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Drop orphaned strip_tool_call_markup import after syncing with main * Tighten supports_tool_passthrough comment * Re-run CI on current main --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- studio/backend/core/inference/llama_cpp.py | 8 +++++++- studio/backend/routes/inference.py | 19 ++++++++++--------- 2 files changed, 17 insertions(+), 10 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 757467a008..b4f8fe1ca6 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -67,7 +67,6 @@ from core.tool_healing import ( _strip_bracket_tag_calls, apply_tool_strip_patterns, strip_outside_think, - strip_tool_call_markup, ) from utils.native_path_leases import child_env_without_native_path_secret from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback @@ -1781,6 +1780,13 @@ class LlamaCppBackend: return False return self._supports_tools + @property + def supports_tool_passthrough(self) -> bool: + # supports_tools is forced off for DiffusionGemma (its agentic loop drops the + # per-step canvas frames), but client passthrough skips that loop, so it uses + # the real _supports_tools. + return self._supports_tools + @property def cache_type_kv(self) -> Optional[str]: return self._cache_type_kv diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 5332037e0d..ce755ae1fb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -6053,14 +6053,13 @@ async def openai_chat_completions( # free-form sampling. Guided decoding does not require ``supports_tools`` -- # the grammar machinery is independent of tool-call parsing. _has_response_format = _extract_response_format(payload) is not None - _tools_passthrough = llama_backend.supports_tools and ( - (payload.tools and len(payload.tools) > 0) or _has_tool_messages - ) - if ( - using_gguf - and not _effective_enable_tools(payload) - and (_tools_passthrough or _has_response_format) - ): + _tools_passthrough = getattr( + llama_backend, "supports_tool_passthrough", llama_backend.supports_tools + ) and ((payload.tools and len(payload.tools) > 0) or _has_tool_messages) + # DiffusionGemma keeps supports_tools off, so the server-side tool loop can't + # claim the request; fall through to client passthrough, matching /v1/messages. + _server_tool_loop = _effective_enable_tools(payload) and llama_backend.supports_tools + if using_gguf and not _server_tool_loop and (_tools_passthrough or _has_response_format): if _wants_multiple_choices(payload): raise _reject_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: @@ -10222,7 +10221,9 @@ async def anthropic_messages( and not _has_image ) client_tools = ( - not server_tools and len(openai_client_tools) > 0 and llama_backend.supports_tools + not server_tools + and len(openai_client_tools) > 0 + and getattr(llama_backend, "supports_tool_passthrough", llama_backend.supports_tools) ) # Anthropic tool_choice.disable_parallel_tool_use caps the response to a From df6b5a57d97ab206dd85572e9760a36ef779622a Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 10:32:06 +0100 Subject: [PATCH 013/402] Fix case-variant model matching and GGUF cache reuse in unsloth start (#6900) * fix: handle case-variant GGUF cache hits for unsloth start * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * gguf cache: keep split shards co-located and isolate cache tests properly When a cached main shard was reused from an older snapshot, the extra shards were resolved independently and could come from a different snapshot dir (or a fresh download into the current ref), leaving llama.cpp unable to load a multi-shard GGUF whose pieces are split across directories. Only reuse a cached main shard when every sibling shard sits in the same snapshot; otherwise fetch the whole set together so they stay co-located. Also patch huggingface_hub.constants.HF_HUB_CACHE (not just the HF_HUB_CACHE env var) in the two cache tests that seeded a temp cache: the snapshot lookup reads the module constant, so the env-only override let the real cache leak in and skip an asserted download. * Do not let a companion-only cache snapshot shadow real GGUF variants When listing GGUF variants from the local HF cache, a newer snapshot may contain only a companion file (for example a vision projector fetched on demand) while the actual quant files live in an older snapshot. The prior scan returned the first snapshot whose vision flag was set, yielding an empty variant list and hiding the real quants. Keep scanning older snapshots for actual variants and carry the vision flag across snapshots. Also record the disk-space fallback variant's size in expected_sizes so the later cache-reuse probe can size-verify the fallback main shard instead of only checking for its existence. * Propagate cached repo casing to companions and preflight split co-location Two fixes to the case-variant GGUF cache reuse: - Resolve the requested repo id to its cached canonical casing once in load_model, up front, and pass it to the main GGUF and its companions (mmproj / MTP drafter). Previously only _download_gguf resolved the casing internally, so a case-variant request loaded the main file from the canonical cache dir while the companions kept the requested casing and missed the cached vision projector / drafter offline. Extracted the resolution into a shared _resolve_repo_id_casing helper. - Apply the split-shard co-location check in the disk-space preflight. When a split GGUF's shards are cached across different snapshots the whole set is refetched later, so counting them as cached made the preflight read 0 bytes to download, skip the smaller-variant fallback, and then fail the full download on a low-disk machine. * Reuse a co-located split GGUF snapshot and fix split fallback size probe - When reusing a cached split GGUF, scan snapshots for one that holds the whole set co-located instead of taking the newest snapshot's first shard. A newer snapshot with only the first shard no longer shadows an older complete snapshot, so an already-cached split model is reused rather than refetched (which would fail offline). - The disk-space fallback records its size in expected_sizes only for a single-file fallback. _find_smallest_fitting_variant returns the whole variant size, so using it as the first shard's expected size rejected a valid cached first shard of a split fallback and forced a re-download. * Scan for a complete split snapshot in the preflight; require a loaded catalog hit - The disk-space preflight now uses the same co-located snapshot scan as the download path (_cached_colocated_split_main) instead of the newest-snapshot probe, so a newer snapshot holding only the first shard no longer masks an older complete one and trips the smaller-variant fallback for a fully cached split model. - _resolve_model only attaches to a /v1/models entry that is actually loaded (loaded != False). /v1/models also lists cached-but-unloaded catalog entries, and matching one by case skipped /api/inference/load and left the agent pointed at a model that is not resident. * Restrict cross-snapshot GGUF cache reuse to offline Reusing a same-name blob from an older or case-variant snapshot bypasses the Hub revision/etag check, so a repo that updates a GGUF in place could serve stale weights online. Gate the cross-snapshot and case-variant reuse (both the disk-space preflight accounting and the download path) on HF_HUB_OFFLINE. Online, hf_hub_download fetches the current revision and resumes a partial download, so the reuse is unnecessary there; offline it remains the resilience fallback. Marked the two reuse regression tests as the offline scenarios they represent and added an online test asserting a fresh fetch. * Harden offline cache reuse and hub-id detection Three follow-ups on the case-variant GGUF cache path: - Honor every truthy HF_HUB_OFFLINE spelling (1/true/yes/on), not just "1", when gating the cross-snapshot and case-variant cache reuse. With HF_HUB_OFFLINE=true the Hub calls are already offline, so the reuse must trigger or the cached GGUF fails to load; route both the preflight accounting and the download path through the same offline parse the rest of the backend uses. - Resolve mmproj/MTP companions from the actual cached snapshot when offline. resolve_cached_repo_id_case can keep a partial lower-case spelling when any dir exists under the requested casing, so an hf_hub_download on that casing misses the canonical companion; scan every case-variant snapshot and return the cached path. - Restrict the case-insensitive model-id match to syntactically valid hub ids (a single namespace/name over the HF charset). A server-side relative path such as models/Llama/Foo.gguf is no longer treated as a hub id, so it cannot casefold-match a differently cased path on a case-sensitive filesystem. This is host independent, unlike the local-existence probe which cannot see a server path. * Only casefold-match model ids against a loopback Studio A two-segment string like Models/Foo is indistinguishable from a hub id, and the local Path.exists() probe in _is_hub_model_id cannot see a path that exists only on a remote Studio host. So against a remote server, casefolding could attach to a distinct server-side path (Models/Foo vs models/foo) on a case-sensitive filesystem. Gate the case-insensitive match on is_loopback_url(base): only a local Studio, where the existence probe is authoritative, casefolds. For a remote Studio the match is exact and a case-mismatched request falls through to /api/inference/load, whose already-loaded dedup resolves it correctly. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: Wasim Yousef Said --- studio/backend/core/inference/llama_cpp.py | 239 +++++++++++- .../tests/test_offline_gguf_cache_fallback.py | 352 +++++++++++++++++- studio/backend/utils/models/model_config.py | 54 ++- unsloth_cli/commands/start.py | 80 +++- unsloth_cli/tests/test_start.py | 183 +++++++++ 5 files changed, 874 insertions(+), 34 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b4f8fe1ca6..b07bd33076 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -377,6 +377,19 @@ def _probe_dns_dead(host: str = "huggingface.co", timeout: float = 2.0) -> bool: return True if result[0] is None else result[0] +def _hf_env_offline() -> bool: + """True when an HF offline env var is set to any truthy value (1/true/yes/on). + + Mirrors utils.models.model_config._env_offline so a user-set HF_HUB_OFFLINE=true + (not just "1") still routes through the local-cache reuse path below. + """ + try: + from utils.models.model_config import _env_offline + return _env_offline() + except Exception: + return os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in {"1", "true", "yes", "on"} + + @contextlib.contextmanager def _hf_offline_if_dns_dead(): """Set HF_HUB_OFFLINE for this block only when DNS to huggingface.co fails; @@ -837,6 +850,112 @@ def _gguf_snapshot_files(snapshot: Path) -> list[str]: ] +def _cached_hf_snapshot_file( + repo_id: str, + filename: str, + *, + expected_size: Optional[int] = None, +) -> Optional[str]: + """Return a cached snapshot file even when HF's current-ref probe misses it.""" + if not filename: + return None + parts = [part for part in filename.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + candidate = snap.joinpath(*parts) + if not candidate.is_file(): + continue + if expected_size: + try: + if candidate.stat().st_size < expected_size: + continue + except OSError: + continue + return str(candidate) + except Exception as e: + logger.debug("Snapshot cache lookup failed for %s/%s: %s", repo_id, filename, e) + return None + + +def _snapshot_has_all_shards( + main_path: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> bool: + """True when every shard sits beside ``main_path`` in the same cache snapshot. + + llama.cpp loads a split GGUF by resolving its siblings from the main shard's + directory, so a cached main shard is only safe to reuse when the rest of the + set is co-located; otherwise the caller must fetch the whole set together. + """ + root = Path(main_path) + for _ in [part for part in main_filename.replace("\\", "/").split("/") if part]: + root = root.parent + for shard in shards: + parts = [part for part in shard.replace("\\", "/").split("/") if part] + if not parts or any(part in (".", "..") for part in parts): + return False + sibling = root.joinpath(*parts) + try: + if not sibling.is_file(): + return False + expected = expected_sizes.get(shard) + if expected and sibling.stat().st_size < expected: + return False + except OSError: + return False + return True + + +def _resolve_repo_id_casing(hf_repo: str) -> str: + """Map a requested repo id to its cached canonical casing, or return it unchanged. + + A case-variant request (for example a lowercased id) resolves to the + canonical-cased cache directory so the main GGUF and its companions + (mmproj / MTP drafter) all read the same cache entry. Returns ``hf_repo`` + unchanged when resolution is unavailable or errors. + """ + try: + from utils.paths import resolve_cached_repo_id_case + return resolve_cached_repo_id_case(hf_repo) + except Exception: + return hf_repo + + +def _cached_colocated_split_main( + repo_id: str, main_filename: str, shards: Iterable[str], expected_sizes: dict[str, int] +) -> Optional[str]: + """Main-shard path from a cache snapshot that also holds every sibling shard. + + A newer snapshot may hold only the first shard while an older snapshot has the + complete split set. ``_cached_hf_snapshot_file`` would return that newer partial + main and the co-location check would then force a refetch, so scan snapshots for + one where the whole set is present and return that main path instead. None when + no snapshot holds the full set. + """ + main_parts = [part for part in main_filename.replace("\\", "/").split("/") if part] + if not main_parts or any(part in (".", "..") for part in main_parts): + return None + try: + from utils.models.model_config import _iter_hf_cache_snapshots + for snap in _iter_hf_cache_snapshots(repo_id): + main_path = snap.joinpath(*main_parts) + if not main_path.is_file(): + continue + expected_main = expected_sizes.get(main_filename) + try: + if expected_main and main_path.stat().st_size < expected_main: + continue + except OSError: + continue + if _snapshot_has_all_shards(str(main_path), main_filename, shards, expected_sizes): + return str(main_path) + except Exception as e: + logger.debug("Co-located split snapshot lookup failed for %s: %s", repo_id, e) + return None + + def _gguf_extra_shards(files: Iterable[str], first_shard: str) -> list[str]: m = _SHARD_FULL_RE.match(first_shard) if not m: @@ -3992,6 +4111,15 @@ class LlamaCppBackend: "Install it with: pip install huggingface_hub" ) + resolved_hf_repo = _resolve_repo_id_casing(hf_repo) + if resolved_hf_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + resolved_hf_repo, + hf_repo, + ) + hf_repo = resolved_hf_repo + # Resolve the filename from the variant gguf_filename = None gguf_extra_shards: list[str] = [] @@ -4037,10 +4165,12 @@ class LlamaCppBackend: # Check disk space; fall back to a smaller variant if needed all_gguf_files = [gguf_filename] + gguf_extra_shards + expected_sizes: dict[str, int] = {} try: from huggingface_hub import get_paths_info, try_to_load_from_cache path_infos = list(get_paths_info(hf_repo, all_gguf_files, token = hf_token)) + expected_sizes = {p.path: p.size for p in path_infos if p.size} total_bytes = sum((p.size or 0) for p in path_infos) # Subtract bytes already in the HF cache so we only preflight @@ -4049,7 +4179,26 @@ class LlamaCppBackend: # cold whenever free disk is below the full weight footprint, # even though nothing needs downloading. already_cached_bytes = 0 - if not force: + # Cross-snapshot / case-variant cache reuse is offline-only (see the download + # path below); online, hf_hub_download fetches the current revision and + # resumes partials, so an old snapshot must not be counted as cached here or + # the preflight would under-count the download and skip the disk fallback. + offline = _hf_env_offline() + # A split GGUF whose shards are not co-located in a single snapshot is + # refetched as a whole set later, so it must not be counted as cached here. + split_needs_refetch = False + if offline and not force and gguf_extra_shards: + # Scan all snapshots for one that holds the whole set co-located, so a + # newer snapshot with only the first shard does not mask an older + # complete one and needlessly trip the disk fallback. + if ( + _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + is None + ): + split_needs_refetch = True + if not force and not split_needs_refetch: for p in path_infos: if not p.size: continue @@ -4057,6 +4206,15 @@ class LlamaCppBackend: cached_path = try_to_load_from_cache(hf_repo, p.path) except Exception: cached_path = None + if ( + not (isinstance(cached_path, str) and os.path.exists(cached_path)) + and offline + ): + cached_path = _cached_hf_snapshot_file( + hf_repo, + p.path, + expected_size = p.size, + ) if isinstance(cached_path, str) and os.path.exists(cached_path): try: on_disk = os.path.getsize(cached_path) @@ -4119,6 +4277,13 @@ class LlamaCppBackend: ) else: gguf_extra_shards = [] + # Record the fallback's size so the later cache-reuse probe can + # size-verify it; only for a single-file fallback, since + # _find_smallest_fitting_variant returns the whole-variant size + # and using that as the first shard's expected size would reject + # a valid cached first shard of a split fallback. + if not gguf_extra_shards: + expected_sizes[fallback_file] = fallback_size else: raise RuntimeError( f"Not enough disk space to download any variant. " @@ -4138,25 +4303,45 @@ class LlamaCppBackend: raise RuntimeError("Cancelled") dl_start = time.monotonic() # Xet primary, HTTP fallback on stall; per-file so finished shards stay cached. - local_path = hf_hub_download_with_xet_fallback( - hf_repo, - gguf_filename, - hf_token, - cancel_event = cancel_event, - on_status = lambda m: logger.info(m), - force_download = force, - ) - for shard in gguf_extra_shards: - if cancel_event.is_set(): - raise RuntimeError("Cancelled") - logger.info(f"Resolving GGUF shard: {shard}") - hf_hub_download_with_xet_fallback( + local_path = None + # Reuse a cached copy from another snapshot / case-variant repo dir only when + # offline. Online, fall through to hf_hub_download so its revision/etag check + # fetches the current file (and resumes a partial) instead of serving a stale + # same-name blob from an older revision. + if not force and _hf_env_offline(): + if gguf_extra_shards: + # A split GGUF must load every shard from one snapshot; reuse only a + # snapshot that holds the whole set co-located, scanning past a newer + # snapshot that has just the first shard while an older one is complete. + local_path = _cached_colocated_split_main( + hf_repo, gguf_filename, gguf_extra_shards, expected_sizes + ) + else: + local_path = _cached_hf_snapshot_file( + hf_repo, + gguf_filename, + expected_size = expected_sizes.get(gguf_filename), + ) + if local_path is None: + local_path = hf_hub_download_with_xet_fallback( hf_repo, - shard, + gguf_filename, hf_token, cancel_event = cancel_event, + on_status = lambda m: logger.info(m), force_download = force, ) + for shard in gguf_extra_shards: + if cancel_event.is_set(): + raise RuntimeError("Cancelled") + logger.info(f"Resolving GGUF shard: {shard}") + hf_hub_download_with_xet_fallback( + hf_repo, + shard, + hf_token, + cancel_event = cancel_event, + force_download = force, + ) except Exception as e: if isinstance(e, RuntimeError) and "Cancelled" in str(e): raise @@ -4234,6 +4419,17 @@ class LlamaCppBackend: if target is None or cancel_event.is_set(): return None + # Offline, resolve the companion straight from the cache snapshot that + # holds it. resolve_cached_repo_id_case can return a partial lower-case + # spelling when any dir exists under the requested casing, so calling + # hf_hub_download with hf_repo would miss the canonical file and silently + # drop the companion. _cached_hf_snapshot_file scans every case variant. + if _hf_env_offline(): + cached = _cached_hf_snapshot_file(hf_repo, target) + if cached: + logger.info("Resolved %s from local HF cache: %s", label, cached) + return cached + try: logger.info(f"Downloading {label}: {hf_repo}/{target}") # Same policy; companions are best-effort (caller below swallows failures to None). @@ -5012,6 +5208,19 @@ class LlamaCppBackend: # dead; cleanup runs even on exception so a transient hiccup # can't quarantine future loads. if hf_repo: + # Resolve the requested repo id to its cached canonical casing once, + # up front, so the main GGUF and its companions (mmproj / MTP drafter) + # all resolve from the same cache entry. Otherwise a case-variant + # request resolves the main file from the canonical cache dir while the + # companions keep the requested casing and miss the cached files. + _resolved_repo = _resolve_repo_id_casing(hf_repo) + if _resolved_repo != hf_repo: + logger.info( + "Using cached repo_id casing '%s' for requested '%s'", + _resolved_repo, + hf_repo, + ) + hf_repo = _resolved_repo with _hf_offline_if_dns_dead(): model_path = self._download_gguf( hf_repo = hf_repo, diff --git a/studio/backend/tests/test_offline_gguf_cache_fallback.py b/studio/backend/tests/test_offline_gguf_cache_fallback.py index 4499881c4d..e24e2ca451 100644 --- a/studio/backend/tests/test_offline_gguf_cache_fallback.py +++ b/studio/backend/tests/test_offline_gguf_cache_fallback.py @@ -79,9 +79,11 @@ from huggingface_hub import constants as hf_constants from core.inference.llama_cpp import ( LlamaCppBackend, + _cached_colocated_split_main, _gguf_files_for_variant, _hf_offline_if_dns_dead, _probe_dns_dead, + _resolve_repo_id_casing, ) from utils.models.model_config import ( _detect_gguf_from_hf_cache, @@ -217,7 +219,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch( "huggingface_hub.list_repo_files", @@ -239,6 +241,214 @@ class TestGgufVariantFileResolution: assert downloaded == ["tinyllamas/stories260K.gguf"] assert out == "/fake/ggml-org/models/tinyllamas/stories260K.gguf" + def test_download_reuses_older_snapshot_when_current_ref_snapshot_is_partial( + self, monkeypatch, hf_cache + ): + # Cross-snapshot reuse is an offline-resilience path: online, hf_hub_download + # resumes the partial current-ref download and revalidates the revision instead + # of serving an older snapshot's same-name blob. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache( + hf_cache, + repo, + {"model-UD-Q4_K_XL.gguf": 4}, + snapshot_sha = "a" * 40, + ) + _build_cache( + hf_cache, + repo, + {"mtp-model.gguf": 1}, + snapshot_sha = "b" * 40, + ) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf", "mtp-model.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_reuses_cached_gguf_when_lowercase_partial_cache_shadows_it( + self, monkeypatch, hf_cache + ): + # Case-variant cross-dir reuse is offline-only; online the canonical repo id + # resolves up front and hf_hub_download fetches the current revision. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + gguf_file = "gemma-4-E2B-it-UD-Q4_K_XL.gguf" + snap = _build_cache( + hf_cache, + canonical_repo, + {gguf_file: 4}, + snapshot_sha = "a" * 40, + ) + lower_snap = _build_cache( + hf_cache, + requested_repo, + {"mtp-gemma-4-E2B-it.gguf": 1}, + snapshot_sha = "b" * 40, + ) + os.utime(lower_snap, (2000, 2000)) + os.utime(snap, (1000, 1000)) + seen_repos: list[str] = [] + + def fake_list_repo_files(repo_id, token = None): + seen_repos.append(repo_id) + return [gguf_file] + + def fake_get_paths_info( + repo_id, + paths, + token = None, + ): + seen_repos.append(repo_id) + return [_types.SimpleNamespace(path = path, size = 4) for path in paths if path] + + def fake_cache(repo_id, filename, *args, **kwargs): + seen_repos.append(repo_id) + return str(snap / filename) if repo_id == canonical_repo else None + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", fake_cache), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf( + hf_repo = requested_repo, + hf_variant = "UD-Q4_K_XL", + ) + + assert out == str(snap / gguf_file) + assert seen_repos + + def test_download_online_does_not_reuse_old_snapshot(self, monkeypatch, hf_cache): + # Online, an older same-name snapshot must not be served (it may be a stale + # revision); hf_hub_download is called so the current revision is fetched and + # its etag revalidated. + monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **kwargs, + ): + downloaded.append(filename) + return f"/fresh/{filename}" + + with ( + patch( + "huggingface_hub.list_repo_files", + lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"], + ), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert downloaded == ["model-UD-Q4_K_XL.gguf"] + assert out == "/fresh/model-UD-Q4_K_XL.gguf" + + def test_download_reuses_older_snapshot_when_offline_env_is_true(self, monkeypatch, hf_cache): + # HF_HUB_OFFLINE accepts truthy spellings beyond "1" (true/yes/on); the offline + # cache reuse must trigger for those too, otherwise the earlier Hub calls run + # offline while this branch still attempts hf_hub_download and the cached GGUF + # cannot load. + monkeypatch.setenv("HF_HUB_OFFLINE", "true") + backend = LlamaCppBackend() + repo = "unsloth/vision-GGUF" + old = _build_cache(hf_cache, repo, {"model-UD-Q4_K_XL.gguf": 4}, snapshot_sha = "a" * 40) + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fail_download(*_args, **_kwargs): + raise AssertionError("should reuse the cached GGUF instead of downloading") + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: ["model-UD-Q4_K_XL.gguf"]), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "UD-Q4_K_XL") + + assert out == str(old / "model-UD-Q4_K_XL.gguf") + + def test_download_companion_resolves_from_case_variant_snapshot_offline( + self, monkeypatch, hf_cache + ): + # Offline, resolve_cached_repo_id_case can keep a partial lower-case spelling, + # so the companion (mmproj) must resolve from whichever case-variant snapshot + # actually holds it rather than being dropped by an hf_hub_download on the + # wrong casing. + monkeypatch.setenv("HF_HUB_OFFLINE", "1") + backend = LlamaCppBackend() + canonical_repo = "unsloth/gemma-4-E2B-it-GGUF" + requested_repo = "unsloth/gemma-4-e2b-it-gguf" + snap = _build_cache(hf_cache, canonical_repo, {"mmproj-F16.gguf": 4}, snapshot_sha = "a" * 40) + # A partial lower-case dir exists so casing resolution keeps the requested spelling. + _build_cache(hf_cache, requested_repo, {"config.json": 1}, snapshot_sha = "b" * 40) + + _offline_exc = type("OfflineModeIsEnabled", (Exception,), {}) + + def fake_list_repo_files(repo_id, token = None): + raise _offline_exc("offline") + + def fail_download(*_args, **_kwargs): + raise AssertionError("should resolve the companion from cache, not download") + + with ( + patch("huggingface_hub.list_repo_files", fake_list_repo_files), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fail_download), + ): + out = backend._download_mmproj(hf_repo = requested_repo) + + assert out == str(snap / "mmproj-F16.gguf") + def test_download_includes_uppercase_split_gguf_shards(self, monkeypatch, tmp_path): backend = LlamaCppBackend() downloaded: list[str] = [] @@ -264,7 +474,7 @@ class TestGgufVariantFileResolution: downloaded.append(filename) return f"/fake/{repo_id}/{filename}" - monkeypatch.setenv("HF_HUB_CACHE", str(tmp_path)) + monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path)) with ( patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), patch("huggingface_hub.get_paths_info", fake_get_paths_info), @@ -279,6 +489,48 @@ class TestGgufVariantFileResolution: assert downloaded == files assert out == "/fake/org/repo/model-Q4_K_M-00001-of-00002.GGUF" + def test_download_refetches_split_gguf_when_shards_span_snapshots(self, monkeypatch, hf_cache): + # The cached main shard lives in an older snapshot; its sibling shard is only + # in a newer, separate snapshot. Reusing the main shard alone would leave + # llama.cpp unable to resolve the sibling, so the whole set must be re-fetched + # together (co-located) rather than served split across snapshot dirs. + backend = LlamaCppBackend() + repo = "org/split" + files = [ + "model-Q4_K_M-00001-of-00002.gguf", + "model-Q4_K_M-00002-of-00002.gguf", + ] + _build_cache(hf_cache, repo, {files[0]: 4}, snapshot_sha = "a" * 40) + _build_cache(hf_cache, repo, {files[1]: 4}, snapshot_sha = "b" * 40) + downloaded: list[str] = [] + + def fake_get_paths_info( + _repo_id, + paths, + token = None, + ): + return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p] + + def fake_download( + repo_id, + filename, + token = None, + **_kwargs, + ): + downloaded.append(filename) + return f"/fake/{repo_id}/{filename}" + + with ( + patch("huggingface_hub.list_repo_files", lambda *_a, **_k: files), + patch("huggingface_hub.get_paths_info", fake_get_paths_info), + patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None), + patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download), + ): + out = backend._download_gguf(hf_repo = repo, hf_variant = "Q4_K_M") + + assert downloaded == files + assert out == f"/fake/{repo}/{files[0]}" + def _siblings(items: dict[str, int]): """Mock ``hf_model_info(...).siblings`` payload.""" @@ -315,6 +567,21 @@ class TestIterHfCacheSnapshots: out = list(_iter_hf_cache_snapshots("unsloth/multi")) assert [p.name for p in out] == ["b" * 40, "a" * 40] + def test_skips_snapshot_when_mtime_is_unavailable(self, hf_cache, monkeypatch): + stale = _build_cache(hf_cache, "unsloth/multi", {"x.gguf": 1}, snapshot_sha = "a" * 40) + good = _build_cache(hf_cache, "unsloth/multi", {"y.gguf": 1}, snapshot_sha = "b" * 40) + original_stat = Path.stat + + def flaky_stat(self, *args, **kwargs): + if self == stale: + raise FileNotFoundError(str(self)) + return original_stat(self, *args, **kwargs) + + monkeypatch.setattr(Path, "stat", flaky_stat) + + out = list(_iter_hf_cache_snapshots("unsloth/multi")) + assert out == [good] + def test_repo_id_match_is_case_insensitive(self, hf_cache): _build_cache(hf_cache, "unsloth/Foo-GGUF", {"Foo-Q4_K_M.gguf": 1}) # Lookup with different org/name casing still resolves @@ -347,6 +614,87 @@ class TestListGgufVariantsFromCache: assert _list_gguf_variants_from_hf_cache("unsloth/absent") is None +class TestCachedColocatedSplitMain: + def test_prefers_older_complete_snapshot_over_newer_partial(self, hf_cache): + # Newer snapshot has only shard 1; older snapshot has the complete set. The + # complete older snapshot must win so the split GGUF can load co-located. + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + old = _build_cache( + hf_cache, "unsloth/split-GGUF", {shard1: 100, shard2: 100}, snapshot_sha = "a" * 40 + ) + new = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "b" * 40) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + main = _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) + assert main is not None + assert main.startswith(str(old)) + + def test_returns_none_when_shards_span_snapshots(self, hf_cache): + shard1 = "m-00001-of-00002.gguf" + shard2 = "m-00002-of-00002.gguf" + a = _build_cache(hf_cache, "unsloth/split-GGUF", {shard1: 100}, snapshot_sha = "a" * 40) + b = _build_cache(hf_cache, "unsloth/split-GGUF", {shard2: 100}, snapshot_sha = "b" * 40) + os.utime(a, (1000, 1000)) + os.utime(b, (2000, 2000)) + + assert _cached_colocated_split_main("unsloth/split-GGUF", shard1, [shard2], {}) is None + + +class TestResolveRepoIdCasing: + def test_maps_to_canonical_casing(self, monkeypatch): + monkeypatch.setattr( + "utils.paths.resolve_cached_repo_id_case", + lambda repo: "unsloth/Gemma-4-GGUF" if repo.lower() == "unsloth/gemma-4-gguf" else repo, + ) + # A companion download passed the resolved id reads the same cache entry + # as the main GGUF instead of missing it under the requested casing. + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/Gemma-4-GGUF" + + def test_passthrough_on_resolver_error(self, monkeypatch): + def boom(_repo): + raise RuntimeError("resolver unavailable") + + monkeypatch.setattr("utils.paths.resolve_cached_repo_id_case", boom) + assert _resolve_repo_id_casing("unsloth/gemma-4-gguf") == "unsloth/gemma-4-gguf" + + def test_companion_only_newer_snapshot_does_not_shadow_real_variants(self, hf_cache): + # A newer snapshot holds only a vision projector fetched on demand, + # while the quant files live in an older snapshot. The newer snapshot + # must not shadow the real variants; the vision flag carries over. + old = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"vision-Q4_K_M.gguf": 100}, + snapshot_sha = "a" * 40, + ) + new = _build_cache( + hf_cache, + "unsloth/vision-GGUF", + {"mmproj-vision-F16.gguf": 10}, + snapshot_sha = "b" * 40, + ) + os.utime(old, (1000, 1000)) + os.utime(new, (2000, 2000)) + + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert [v.quant for v in variants] == ["Q4_K_M"] + assert has_vision is True + + def test_companion_only_cache_returns_empty_variants_with_vision(self, hf_cache): + # Only a vision projector is cached anywhere: report the vision flag + # with an empty variant list rather than None. + _build_cache(hf_cache, "unsloth/vision-GGUF", {"mmproj-vision-F16.gguf": 10}) + out = _list_gguf_variants_from_hf_cache("unsloth/vision-GGUF") + assert out is not None + variants, has_vision = out + assert variants == [] + assert has_vision is True + + class TestListGgufVariantsOffline: def test_offline_env_short_circuits_api(self, hf_cache, clean_offline_env, monkeypatch): _build_cache(hf_cache, "unsloth/a", {"a-UD-Q4_K_XL.gguf": 1}) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5d8458e5f0..281ca24281 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -1617,36 +1617,60 @@ def _iter_hf_cache_snapshots(repo_id: str): cache_dir = Path(hf_constants.HF_HUB_CACHE) target = f"models--{repo_id.replace('/', '--')}".lower() - repo_dir: Optional[Path] = None + repo_dirs: list[Path] = [] try: if not cache_dir.is_dir(): return for entry in cache_dir.iterdir(): if entry.is_dir() and entry.name.lower() == target: - repo_dir = entry - break + repo_dirs.append(entry) except OSError: return - if repo_dir is None: + if not repo_dirs: return - snapshots = repo_dir / "snapshots" - try: - if not snapshots.is_dir(): - return - snap_dirs = [s for s in snapshots.iterdir() if s.is_dir()] - except OSError: + snap_dirs: list[Path] = [] + for repo_dir in repo_dirs: + snapshots = repo_dir / "snapshots" + try: + if snapshots.is_dir(): + for snap_dir in snapshots.iterdir(): + try: + if snap_dir.is_dir(): + snap_dirs.append(snap_dir) + except OSError: + continue + except OSError: + continue + if not snap_dirs: return - snap_dirs.sort(key = lambda s: s.stat().st_mtime, reverse = True) - yield from snap_dirs + snap_dirs_with_mtime = [] + for snap_dir in snap_dirs: + try: + snap_dirs_with_mtime.append((snap_dir.stat().st_mtime, snap_dir)) + except OSError: + continue + snap_dirs_with_mtime.sort(key = lambda item: item[0], reverse = True) + yield from (snap_dir for _, snap_dir in snap_dirs_with_mtime) def _list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]: - """Variants from the local HF cache snapshot, or None if not cached.""" + """Variants from the local HF cache snapshot, or None if not cached. + + A newer snapshot can hold only a companion file (for example a vision + projector fetched on demand) while the quant files live in an older + snapshot. Returning the first snapshot that merely reports a vision flag + would shadow those real variants, so keep scanning older snapshots for + actual variants and carry the vision flag across snapshots. + """ + any_vision = False for snap in _iter_hf_cache_snapshots(repo_id): variants, has_vision = list_local_gguf_variants(str(snap)) - if variants or has_vision: - return variants, has_vision + any_vision = any_vision or has_vision + if variants: + return variants, any_vision + if any_vision: + return [], True return None diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 01014de4d3..477a47cc3d 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -597,6 +597,59 @@ def _loaded_models(base: str, key: str) -> list: return _http_json("GET", f"{base}/v1/models", key, error = "Couldn't list models").get("data", []) +_HF_REPO_ID_SEGMENT_RE = re.compile(r"^[A-Za-z0-9._-]+$") + + +def _is_hub_model_id(value: object) -> bool: + if not isinstance(value, str): + return False + text = value.strip() + if "\\" in text: + return False + if text.startswith(("/", "./", "../", "~")): + return False + if len(text) >= 2 and text[1] == ":" and text[0].isalpha(): + return False + # A hub id is exactly "namespace/name" over a restricted charset. Anything with + # extra path segments (e.g. a server-side relative path such as + # models/Llama/Foo.gguf on a remote Studio) is not a hub id and must not be + # casefold-matched against a differently cased path on a case-sensitive + # filesystem. This is host independent, unlike the existence probe below which + # cannot see a path that only exists on the server. + parts = text.split("/") + if len(parts) != 2: + return False + if any(part in ("", ".", "..") or not _HF_REPO_ID_SEGMENT_RE.match(part) for part in parts): + return False + try: + if Path(os.path.expanduser(text)).exists(): + return False + except OSError: + return False + return True + + +def _model_id_matches( + actual: object, + requested: object, + *, + allow_casefold: bool = True, +) -> bool: + if actual == requested: + return True + # Case-insensitive matching is only safe when the local existence probe in + # _is_hub_model_id is authoritative, i.e. against a loopback Studio on this host. + # Against a remote Studio a two-segment string is indistinguishable from a + # server-side relative path (e.g. Models/Foo vs models/foo), so casefolding it + # could attach to the wrong model on a case-sensitive server; defer to an exact + # match there and let the load endpoint resolve the requested path. + if not allow_casefold: + return False + if not (_is_hub_model_id(actual) and _is_hub_model_id(requested)): + return False + return str(actual).casefold() == str(requested).casefold() + + def _resolve_model( base: str, key: str, @@ -604,6 +657,9 @@ def _resolve_model( load: LoadOptions = LoadOptions(), ) -> dict: models = _loaded_models(base, key) + # Only casefold-match ids against a loopback Studio, where _is_hub_model_id's + # local existence probe can actually reject a server-side path; see the note there. + allow_casefold = is_loopback_url(base) # /v1/models reports the model id but not the active GGUF variant or runtime load # settings, so an id match alone can hide the wrong quant (Q8_0 serving while the # user asked for UD-Q4_K_XL). When the user passed any explicit load knob, defer to @@ -613,10 +669,21 @@ def _resolve_model( load_has_overrides = bool( load.gguf_variant or load.max_seq_length or not load.load_in_4bit or load.tensor_parallel ) + # /v1/models also lists cached-but-unloaded catalog entries (loaded == False); + # matching one would skip /api/inference/load and leave the agent pointed at a + # model that is not resident, so only attach to an entry that is actually loaded. match = ( None if requested and load_has_overrides - else next((m for m in models if m["id"] == requested), None) + else next( + ( + m + for m in models + if _model_id_matches(m.get("id"), requested, allow_casefold = allow_casefold) + and m.get("loaded") is not False + ), + None, + ) ) if requested and match is None: typer.echo( @@ -651,7 +718,16 @@ def _resolve_model( if isinstance(loaded, dict): wanted |= {loaded.get("model"), loaded.get("display_name")} - {None} models = _loaded_models(base, key) - match = next((m for m in models if m["id"] in wanted), None) + match = next( + ( + m + for m in models + if any( + _model_id_matches(m.get("id"), w, allow_casefold = allow_casefold) for w in wanted + ) + ), + None, + ) if match is not None: return match if requested: diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index bd964d5e54..ee5b442c27 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -457,6 +457,189 @@ def test_connect_codex_no_launch(fake_studio, tmp_path): assert (home / "unsloth_api.config.toml").exists() +def test_connect_codex_matches_requested_model_case_insensitively(fake_studio, tmp_path): + result = CliRunner().invoke( + start.start_app, + [ + "codex", + "--no-launch", + "--model", + "unsloth/gemma-4-26b-a4b-it-gguf", + ], + ) + assert result.exit_code == 0, result.output + home = tmp_path / "agents" / "codex" + profile = _parse_toml((home / "unsloth_api.config.toml").read_text()) + assert profile["model"] == MODEL["id"] + + +def test_resolve_model_matches_loaded_canonical_case_after_load(monkeypatch): + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url, payload)) + if url.endswith("/v1/models"): + return { + "data": [ + { + "id": "unsloth/gemma-4-E2B-it-GGUF" if state["loaded"] else "other/model", + "context_length": 131072, + } + ] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/gemma-4-E2B-it-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model( + BASE, + "sk-test", + "unsloth/gemma-4-e2b-it-gguf", + start.LoadOptions(gguf_variant = "UD-Q4_K_XL"), + ) + + assert entry["id"] == "unsloth/gemma-4-E2B-it-GGUF" + assert any(c[1].endswith("/api/inference/load") for c in calls) + + +def test_resolve_model_loads_when_catalog_hit_is_not_loaded(monkeypatch): + # A cached-but-unloaded catalog entry (loaded == False) that only case-differs must + # not be treated as ready; the load endpoint must still be called so the requested + # model becomes resident instead of the agent preflighting a different backend. + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [ + { + "id": "unsloth/Gemma-4-GGUF", + "loaded": state["loaded"], + "context_length": 131072, + } + ] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/Gemma-4-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf") + + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_resolve_model_attaches_to_loaded_catalog_hit_without_reload(monkeypatch): + # The mirror case: a loaded entry (loaded == True) that case-matches attaches with + # no /api/inference/load call. + calls = [] + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}] + } + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model(BASE, "sk-test", "unsloth/gemma-4-gguf") + + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert not any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_resolve_model_remote_studio_does_not_casefold_attach(monkeypatch): + # Against a remote Studio the local existence probe cannot see server-side paths, + # so a case-variant loaded id must NOT attach without a load: it could be a distinct + # server-side path on a case-sensitive host. The load endpoint resolves the request. + calls = [] + state = {"loaded": False} + + def http_json( + method, + url, + token, + payload = None, + timeout = 30, + error = None, + ): + calls.append((method, url)) + if url.endswith("/v1/models"): + return { + "data": [{"id": "unsloth/Gemma-4-GGUF", "loaded": True, "context_length": 131072}] + } + if url.endswith("/api/inference/load"): + state["loaded"] = True + return {"model": "unsloth/Gemma-4-GGUF"} + raise AssertionError(f"unexpected request: {method} {url}") + + monkeypatch.setattr(start, "_http_json", http_json) + + entry = start._resolve_model("http://10.0.0.5:8888", "sk-test", "unsloth/gemma-4-gguf") + + # The load endpoint was consulted (no casefold shortcut), and we still attach to the + # server's canonical id it reports back. + assert entry["id"] == "unsloth/Gemma-4-GGUF" + assert any(u.endswith("/api/inference/load") for _, u in calls) + + +def test_model_id_matching_does_not_casefold_local_paths(tmp_path): + existing_local = tmp_path / "Org" / "Foo" + existing_local.mkdir(parents = True) + + assert start._model_id_matches("Org/Foo", "org/foo") + assert not start._model_id_matches(str(existing_local), str(existing_local).lower()) + assert not start._model_id_matches("./Models/Foo", "./models/foo") + assert not start._model_id_matches(r".\Models\Foo", r".\models\foo") + # A server-side relative path (extra path segments) is not a hub id even when it + # does not exist on the CLI host, so it must not casefold-match a differently + # cased path on a case-sensitive server filesystem. + assert not start._is_hub_model_id("models/Llama/Foo.gguf") + assert not start._model_id_matches("models/Llama/Foo.gguf", "models/llama/foo.gguf") + # A genuine two-segment hub id still matches case-insensitively. + assert start._is_hub_model_id("unsloth/Gemma-3-4b-it-GGUF") + assert start._model_id_matches("unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf") + # Casefolding is gated to loopback studios (allow_casefold). With it disabled (a + # remote studio, where a two-segment string could be a server-side path), even a + # genuine hub-id case variant must not match, so the load endpoint resolves it. + assert not start._model_id_matches( + "unsloth/Gemma-3-4b-it-GGUF", "unsloth/gemma-3-4b-it-gguf", allow_casefold = False + ) + assert start._model_id_matches("unsloth/Foo", "unsloth/Foo", allow_casefold = False) + + def test_connect_codex_launch_uses_ephemeral_home(fake_studio, monkeypatch): # Launch mode writes config to a throwaway temp CODEX_HOME and removes it after # the agent exits; the user's real ~/.codex is never the target. From f1a2621631d3adfd72a65dc9710290abb81c54db Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 03:09:52 -0700 Subject: [PATCH 014/402] Studio: show Hugging Face address on hover for Hub and online model rows (#6382) (#6928) * Studio: show Hugging Face address on hover for Hub and online model rows The model selector already shows an on-disk path tooltip on local rows, but Hub and online rows showed only the bare repo id, and nothing at all when there was no VRAM estimate. Add an optional hubUrl prop and a hubRepoUrl helper that mirrors localPathTooltip, and surface huggingface.co/ on hover for the Discover, search, and downloaded Hub rows. Local and VRAM tooltips are unchanged; the VRAM tooltip now also appends the address line. Closes #6382 * Studio: use a 700ms hover delay before the model-row tooltip Give the model-row hover tooltip (the Hugging Face address, plus the VRAM and local-path lines it shares) a 700ms open delay instead of showing it instantly, so it does not flash while sweeping the mouse down the list. * Fix/adjust GGUF tooltips for PR #6928 --------- Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: Wasim Yousef Said --- .../assistant-ui/model-selector/pickers.tsx | 62 ++++++++++++++----- 1 file changed, 45 insertions(+), 17 deletions(-) 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 9890c5f574..ff0351883d 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -393,6 +393,7 @@ function ModelRow({ vramEst, gpuGb, tooltipText, + hubUrl, optionProps, onArrowDownIntoChildren, capabilities, @@ -409,6 +410,10 @@ function ModelRow({ vramEst?: number; gpuGb?: number; tooltipText?: ReactNode; + /** Hugging Face address (e.g. "huggingface.co/owner/name") for online/Hub + * rows; surfaced on hover so their repo id / URL is discoverable the same + * way local rows show an on-disk path. Omit to show no address line. */ + hubUrl?: string; optionProps?: ModelRowOptionProps; onArrowDownIntoChildren?: () => boolean; /** Capability override (HF rows have tags); falls back to name detection. */ @@ -546,30 +551,41 @@ function ModelRow({ ); - if (vramTooltipText) { - return ( - - {content} - - {label} - {vramTooltipText} - - - ); - } + // Optional Hugging Face address line for online/Hub rows, rendered under + // whichever tooltip shows so the repo id / URL is always visible on hover. + const hubUrlLine = hubUrl ? ( + + {hubUrl} + + ) : null; - if (tooltipText) { + const tooltipBody = vramTooltipText ? ( + <> + {label} + {vramTooltipText} + {hubUrlLine} + + ) : tooltipText ? ( + <> + {tooltipText} + {hubUrlLine} + + ) : hubUrl ? ( + <> + {label} + {hubUrlLine} + + ) : null; + + if (tooltipBody) { return ( - + {content} - {tooltipText} + {tooltipBody} ); @@ -1193,6 +1209,13 @@ function localPathTooltip(name: string, path: string): ReactNode { ); } +/** Hugging Face address for an online/Hub row, or undefined when the repo id is + * missing so the row shows no (empty) address line on hover. */ +function hubRepoUrl(id: string | null | undefined): string | undefined { + const trimmed = id?.trim(); + return trimmed ? `huggingface.co/${trimmed}` : undefined; +} + /** Whether a local model is an MLX build (name hint). MLX runs on Mac only, so * callers gate visibility on the host being a Mac. */ function localModelIsMlx(m: LocalModelInfo): boolean { @@ -2462,6 +2485,7 @@ export function HubModelPicker({
Date: Wed, 8 Jul 2026 03:13:32 -0700 Subject: [PATCH 015/402] Studio: fix currency and indentation edge cases in LaTeX rendering (#6957) * Studio: fix link, currency and indentation edge cases in LaTeX rendering Follow-up to #6914. Three fixes to studio/frontend/src/lib/latex.ts: - Skip reference-link definition URLs ([id]: url) during delimiter conversion, so escaped parens in such URLs are not rewritten as math. - Preserve the opener line's indentation when emitting a display $$ block, so a \[...\] inside a list item stays part of the list. - Stop a currency amount from pairing with a converted span's opening $, which swallowed the price into math (for example $5 + x \(y\)). * Exclude GFM footnote definitions from the reference-URL skip A footnote definition like [^1]: \(x\) had its body treated as a link destination, so leading math was left literal. Skip [^...] labels. * Merge overlapping link destination regions A reference-def token can nest inline-link spans (for example [1]: http://h/[a](b)/foo\(x\)), so the combined spans could overlap and isInRegion's binary search missed the outer one, rewriting the URL. Merge overlapping spans before the search. * Guard lineStart when the display opener is at index 0 Behavior is unchanged (lastIndexOf clamps a negative fromIndex to 0), but the explicit guard avoids relying on that implicit clamp. * Scope to indentation and currency fixes Drop the reference-link URL protection added earlier. It guards a case models effectively never emit (escaped parens in a reference-style URL), and approximating CommonMark reference definitions with a regex needs open-ended special-casing. Keep the two high-value fixes: preserve display math indentation (including multi-line bodies) inside a list item, and stop a currency amount from pairing with a converted span's opening dollar sign. --- studio/frontend/src/lib/latex.ts | 28 +++++++++++++++++++++++++--- 1 file changed, 25 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/lib/latex.ts b/studio/frontend/src/lib/latex.ts index 86a9634048..edf9875602 100644 --- a/studio/frontend/src/lib/latex.ts +++ b/studio/frontend/src/lib/latex.ts @@ -173,7 +173,11 @@ function looksLikeMathBody(body: string): boolean { * (`**$X$**`, `__$X$__`) are always math: LLMs use that for "bold math" * and the heuristic would otherwise reject prose-shaped bodies like "90 - x". */ -function hasInlineMathCloser(content: string, offset: number): boolean { +function hasInlineMathCloser( + content: string, + offset: number, + mathRegions: Array<[number, number]>, +): boolean { const MAX_SPAN = 200; const limit = Math.min(content.length, offset + 1 + MAX_SPAN); for (let i = offset + 1; i < limit; i++) { @@ -181,6 +185,9 @@ function hasInlineMathCloser(content: string, offset: number): boolean { if (c === "\n") return false; if (c !== "$") continue; if (content[i - 1] === "\\") continue; + // A `$` opening a generated span (from `\(...\)`) is not a currency closer; + // pairing with it would swallow the price into math (`$5 + x \(y\)`). + if (isInRegion(i, mathRegions)) return false; if (content[i + 1] === "$") { i++; continue; @@ -294,7 +301,22 @@ function convertLatexDelimiters(content: string): { continue; } append(content.slice(last, match.index)); - const wrapped = isDisplay ? `\n$$\n${body}\n$$\n` : `$${body}$`; + let wrapped: string; + if (isDisplay) { + // Keep the opener's leading indentation so a `$$` block inside a list item + // stays in the container instead of breaking out at column 0. Only when the + // opener is whitespace-prefixed, so inline `text \[x\]` keeps column 0. + const lineStart = + match.index > 0 ? content.lastIndexOf("\n", match.index - 1) + 1 : 0; + const prefix = content.slice(lineStart, match.index); + const indent = /^\s*$/.test(prefix) ? prefix : ""; + // Indent every body line, not just the first, so multi-line display math + // (`\[a\nb\]`) stays wholly inside the container. + const inner = indent ? body.replace(/\n/g, `\n${indent}`) : body; + wrapped = `\n${indent}$$\n${indent}${inner}\n${indent}$$\n`; + } else { + wrapped = `$${body}$`; + } const start = append(wrapped); mathRegions.push([start, offset]); last = matchEnd; @@ -334,7 +356,7 @@ export function preprocessLaTeX(content: string): string { if (isInRegion(offset, mathRegions)) { return match; } - if (hasInlineMathCloser(text, offset)) { + if (hasInlineMathCloser(text, offset, mathRegions)) { return match; } return "\\" + match; From 38dacb8a1f49906d2fa23e4dace198a9c56b5220 Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:26 +0800 Subject: [PATCH 016/402] Add MLX backend support for CLI unsloth train (#6709) * feat(studio): route CLI trainer to MLX backend * fix(studio): harden MLX trainer routing * fix(studio): harden MLX trainer adapter routing * test(studio): assert MLX CLI activation order * fix(studio): address MLX CLI review feedback * feat(cli): support MLX in legacy script * fix(cli): adapt MLX tokenizer for raw text * fix(cli): omit unsupported MLX eval batch arg * fix(cli): feed raw text to MLX trainer * Fix CLI MLX routing and Python 3.9 annotations Route the MLX backend through create_mlx_trainer_adapter so the torch-free Apple Silicon path never imports trainer.py (torch/unsloth/trl). Replace from __future__ import annotations with typing.Optional/Union so the CLI annotations stay Python 3.9 compatible without the unused-import lint hit. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Strip return_tensors from MLX raw-text tokenizer proxy On a torch-free MLX install, RawTextDataLoader calls the tokenizer with return_tensors='pt'; the callable proxy forwarded that to the HF tokenizer, which tried to build torch tensors and failed before training. Drop return_tensors so the MLX path returns plain token ids. * Tighten CLI MLX-backend comments --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/core/training/trainer.py | 31 +- studio/backend/core/training/training.py | 624 +++++++++++++++--- studio/backend/core/training/worker.py | 248 +++++-- .../backend/tests/test_training_preflight.py | 232 +++++++ unsloth-cli.py | 243 ++++--- unsloth_cli/commands/train.py | 46 +- 6 files changed, 1162 insertions(+), 262 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 20b2305a5a..958f8f4197 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -62,7 +62,6 @@ from loggers import get_logger import time from pathlib import Path from typing import Any, Dict, List, Optional, Callable -from dataclasses import dataclass import pandas as pd from datasets import Dataset from utils.datasets.cache_safe import load_dataset_cache_safe as load_dataset @@ -86,6 +85,11 @@ from utils.native_path_leases import child_env_without_native_path_secret from utils.subprocess_compat import ( windows_hidden_subprocess_kwargs as _windows_hidden_subprocess_kwargs, ) +from .training import ( + TrainingProgress, + create_mlx_trainer_adapter, + should_use_mlx_training_backend, +) logger = get_logger(__name__) @@ -104,31 +108,16 @@ def _build_report_targets(training_args) -> list[str] | str: return report_to or "none" -@dataclass -class TrainingProgress: - """Training progress tracking""" - - epoch: float = 0 - step: int = 0 - total_steps: int = 0 - loss: Optional[float] = None - learning_rate: Optional[float] = None - is_training: bool = False - is_completed: bool = False - error: Optional[str] = None - status_message: str = "Ready to train" # Current stage - elapsed_seconds: Optional[float] = None - eta_seconds: Optional[float] = None - grad_norm: Optional[float] = None - num_tokens: Optional[int] = None - eval_loss: Optional[float] = None - - class UnslothTrainer: """ Unsloth Training Backend """ + def __new__(cls, *args, **kwargs): + if cls is UnslothTrainer and should_use_mlx_training_backend(): + return create_mlx_trainer_adapter(*args, **kwargs) + return super().__new__(cls) + def __init__(self): self.model = None self.tokenizer = None diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index f4233fcf04..2ddda19951 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -14,17 +14,19 @@ import json as _json import math import multiprocessing as mp import os +import platform import queue import re import shutil import threading import time +import traceback import structlog from datetime import datetime, timezone from loggers import get_logger -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from pathlib import Path -from typing import Optional, Tuple, Any, TYPE_CHECKING +from typing import Optional, Tuple, Any, Callable, Union, TYPE_CHECKING if TYPE_CHECKING: import matplotlib.pyplot as plt @@ -98,6 +100,107 @@ def _coerce_optional_nonneg_float(name: str, value): return coerced +def is_apple_silicon_training_platform() -> bool: + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def is_mlx_training_device(device: Any) -> bool: + return ( + str(device).lower() == "mlx" + or str(device).lower().endswith(".mlx") + or getattr(device, "name", "").lower() == "mlx" + ) + + +def should_use_mlx_training_backend(*, device: Optional[Any] = None) -> bool: + if device is not None: + return is_mlx_training_device(device) + return is_apple_silicon_training_platform() + + +def _build_training_worker_config(values: dict[str, Any]) -> dict[str, Any]: + """Build the normalized worker config shared by Studio and the CLI adapter.""" + config = { + "model_name": values["model_name"], + "project_name": values.get("project_name"), + "training_type": values.get("training_type", "LoRA/QLoRA"), + "hf_token": values.get("hf_token", ""), + "load_in_4bit": values.get("load_in_4bit", True), + "max_seq_length": values.get("max_seq_length", 2048), + "vision_image_size": values.get("vision_image_size"), + "hf_dataset": values.get("hf_dataset", ""), + "local_datasets": values.get("local_datasets"), + "local_eval_datasets": values.get("local_eval_datasets"), + "format_type": values.get("format_type", ""), + "subset": values.get("subset"), + "train_split": values.get("train_split", "train"), + "eval_split": values.get("eval_split"), + "eval_steps": values.get("eval_steps", 0.00), + "dataset_streaming": values.get("dataset_streaming", False), + "dataset_slice_start": values.get("dataset_slice_start"), + "dataset_slice_end": values.get("dataset_slice_end"), + "custom_format_mapping": values.get("custom_format_mapping"), + "is_dataset_image": values.get("is_dataset_image", False), + "is_dataset_audio": values.get("is_dataset_audio", False), + "is_embedding": values.get("is_embedding", False), + "num_epochs": values.get("num_epochs", 3), + "learning_rate": values.get("learning_rate", "2e-4"), + "embedding_learning_rate": values.get("embedding_learning_rate"), + "batch_size": values.get("batch_size", 2), + "gradient_accumulation_steps": values.get("gradient_accumulation_steps", 4), + "warmup_steps": values.get("warmup_steps"), + "warmup_ratio": values.get("warmup_ratio"), + "max_steps": values.get("max_steps", 0), + "save_steps": values.get("save_steps", 0), + "weight_decay": values.get("weight_decay", 0.001), + "max_grad_norm": values.get("max_grad_norm", 0.0), + "max_grad_value": _coerce_optional_nonneg_float( + "max_grad_value", values.get("max_grad_value") + ), + "max_grad_leaf_norm": _coerce_optional_nonneg_float( + "max_grad_leaf_norm", values.get("max_grad_leaf_norm") + ), + "cast_norm_output_to_input_dtype": _coerce_optional_bool( + values.get("cast_norm_output_to_input_dtype"), True + ), + "random_seed": _coerce_seed(values.get("random_seed")), + "packing": values.get("packing", False), + "optim": values.get("optim", "adamw_8bit"), + "lr_scheduler_type": values.get("lr_scheduler_type", "linear"), + "use_lora": values.get("use_lora", True), + "lora_r": values.get("lora_r", 16), + "lora_alpha": values.get("lora_alpha", 16), + "lora_dropout": values.get("lora_dropout", 0.0), + "target_modules": values.get("target_modules"), + "gradient_checkpointing": values.get("gradient_checkpointing", "unsloth"), + "use_rslora": values.get("use_rslora", False), + "use_loftq": values.get("use_loftq", False), + "train_on_completions": values.get("train_on_completions", False), + "finetune_vision_layers": values.get("finetune_vision_layers", True), + "finetune_language_layers": values.get("finetune_language_layers", True), + "finetune_attention_modules": values.get("finetune_attention_modules", True), + "finetune_mlp_modules": values.get("finetune_mlp_modules", True), + "enable_wandb": values.get("enable_wandb", False), + "wandb_token": values.get("wandb_token"), + "wandb_project": values.get("wandb_project", "unsloth-training"), + "enable_tensorboard": values.get("enable_tensorboard", False), + "tensorboard_dir": values.get("tensorboard_dir", "runs"), + "resume_from_checkpoint": values.get("resume_from_checkpoint"), + "trust_remote_code": values.get("trust_remote_code", False), + "approved_remote_code_fingerprint": values.get("approved_remote_code_fingerprint"), + "subject": values.get("subject"), + "gpu_ids": values.get("gpu_ids"), + "s3_config": values.get("s3_config"), + "disable_xet": values.get("disable_xet", False), + } + for key in ("output_dir", "allow_external_output_dir"): + if key in values: + config[key] = values.get(key) + if config["training_type"] == "Full Finetuning": + config["load_in_4bit"] = False + return config + + _HF_TMP_CHECKPOINT_RE = re.compile(r"^tmp-checkpoint-\d+$") @@ -133,7 +236,7 @@ def _s3_dataset_name(s3_dataset: Any) -> Optional[str]: return f"s3://{bucket}/{prefix}" if prefix else f"s3://{bucket}" -def _cleanup_cancelled_checkpoints(output_dir: str | os.PathLike) -> None: +def _cleanup_cancelled_checkpoints(output_dir: Union[str, os.PathLike]) -> None: """Remove only HF Trainer ``tmp-checkpoint-/`` partials after a cancel. Completed ``checkpoint-/`` dirs survive. Symlinked output_dir / children @@ -183,7 +286,7 @@ PLOT_HEIGHT = 3.5 @dataclass class TrainingProgress: - """Mirror of trainer.TrainingProgress so the parent never imports heavy ML modules.""" + """Shared training progress payload for Studio and backend-aware trainers.""" epoch: float = 0 step: int = 0 @@ -200,6 +303,423 @@ class TrainingProgress: num_tokens: Optional[int] = None eval_loss: Optional[float] = None peak_memory_gb: Optional[float] = None + output_dir: Optional[str] = None + + +class _MLXTrainerAdapter: + """Adapts the legacy UnslothTrainer API to the shared Studio MLX worker path.""" + + def __init__(self): + self.model = None + self.tokenizer = None + self.trainer = None + self.training_thread = None + self.training_progress = TrainingProgress() + self.progress_callbacks: list[Callable[[TrainingProgress], None]] = [] + self.is_training = False + self.should_stop = False + self.save_on_stop = True + self.load_in_4bit = True + self.output_dir = None + + self.is_cpt = False + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model_name = None + self.max_seq_length = None + + self._model_config: dict[str, Any] = {} + self._peft_config: dict[str, Any] = {} + self._dataset_config: dict[str, Any] = {} + self._event_queue: Optional[queue.Queue] = None + self._stop_queue: Optional[queue.Queue] = None + self._pump_thread: Optional[threading.Thread] = None + self._lock = threading.Lock() + + def _activate_transformers_for_model(self, model_name: str, hf_token: Optional[str]) -> None: + try: + from utils.transformers_version import activate_transformers_for_subprocess + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + logger.warning("MLX trainer adapter Transformers activation failed", error = str(exc)) + + def add_progress_callback(self, callback: Callable[[TrainingProgress], None]): + self.progress_callbacks.append(callback) + + def _update_progress(self, **kwargs): + with self._lock: + for key, value in kwargs.items(): + if hasattr(self.training_progress, key): + setattr(self.training_progress, key, value) + progress = self.training_progress + for callback in self.progress_callbacks: + try: + callback(progress) + except Exception: + pass + + def load_model( + self, + model_name: str, + max_seq_length: int = 2048, + load_in_4bit: bool = True, + hf_token: Optional[str] = None, + is_dataset_image: bool = False, + is_dataset_audio: bool = False, + trust_remote_code: bool = False, + full_finetuning: bool = False, + gpu_ids: Optional[list[int]] = None, + ) -> bool: + self.model_name = model_name + self.max_seq_length = max_seq_length + self.load_in_4bit = load_in_4bit + self._audio_type = None + self._activate_transformers_for_model(model_name, hf_token) + try: + from utils.models import detect_audio_type, is_vision_model + + self._audio_type = detect_audio_type(model_name, hf_token) + if self._audio_type == "audio_vlm": + self.is_audio = False + self.is_audio_vlm = bool(is_dataset_audio) + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False + vision = is_vision_model(model_name, hf_token = hf_token) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and bool(is_dataset_image) + except Exception as exc: + logger.warning("MLX trainer adapter model type detection failed", error = str(exc)) + self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False + self.model = object() + self.tokenizer = object() + self._model_config = { + "model_name": model_name, + "max_seq_length": max_seq_length, + "load_in_4bit": load_in_4bit, + "hf_token": hf_token or "", + "is_dataset_image": bool(is_dataset_image), + "is_dataset_audio": bool(is_dataset_audio), + "trust_remote_code": bool(trust_remote_code), + "gpu_ids": gpu_ids, + } + self._update_progress( + is_training = False, + is_completed = False, + error = None, + step = 0, + loss = 0.0, + epoch = 0, + status_message = f"Queued MLX model load: {model_name}", + ) + return True + + def prepare_model_for_training( + self, + use_lora: bool = True, + finetune_vision_layers: bool = True, + finetune_language_layers: bool = True, + finetune_attention_modules: bool = True, + finetune_mlp_modules: bool = True, + target_modules: Optional[Union[list, str]] = None, + lora_r: int = 16, + lora_alpha: int = 16, + lora_dropout: float = 0.0, + use_gradient_checkpointing: Union[str, bool] = "unsloth", + use_rslora: bool = False, + use_loftq: bool = False, + ) -> bool: + self._peft_config = { + "use_lora": bool(use_lora), + "lora_r": lora_r, + "lora_alpha": lora_alpha, + "lora_dropout": lora_dropout, + "target_modules": target_modules, + "gradient_checkpointing": use_gradient_checkpointing, + "use_rslora": bool(use_rslora), + "use_loftq": bool(use_loftq), + "finetune_vision_layers": bool(finetune_vision_layers), + "finetune_language_layers": bool(finetune_language_layers), + "finetune_attention_modules": bool(finetune_attention_modules), + "finetune_mlp_modules": bool(finetune_mlp_modules), + } + self._update_progress(status_message = "Queued MLX training setup") + return True + + def load_and_format_dataset( + self, + dataset_source: Optional[str], + format_type: str = "auto", + local_datasets: Optional[list[str]] = None, + local_eval_datasets: Optional[list[str]] = None, + custom_format_mapping: Optional[dict[str, Any]] = None, + subset: Optional[str] = None, + train_split: str = "train", + eval_split: Optional[str] = None, + dataset_streaming: bool = False, + eval_steps: float = 0.00, + dataset_slice_start: Optional[int] = None, + dataset_slice_end: Optional[int] = None, + is_cpt: bool = False, + s3_config: dict = None, + ) -> Optional[tuple]: + self._dataset_config = { + "hf_dataset": dataset_source or "", + "local_datasets": local_datasets, + "local_eval_datasets": local_eval_datasets, + "format_type": format_type or "", + "custom_format_mapping": custom_format_mapping, + "subset": subset, + "train_split": train_split or "train", + "eval_split": eval_split, + "dataset_streaming": bool(dataset_streaming), + "eval_steps": eval_steps or 0.0, + "dataset_slice_start": dataset_slice_start, + "dataset_slice_end": dataset_slice_end, + "s3_config": s3_config, + } + self.is_cpt = bool(is_cpt) + self._update_progress(status_message = "Queued MLX dataset load") + return ({"dataset": [], "final_format": "deferred_mlx_cli", "success": True}, None) + + def start_training( + self, + dataset = None, + eval_dataset = None, + **training_args, + ) -> bool: + if self.is_training and self.training_thread and self.training_thread.is_alive(): + return False + if self._pump_thread and self._pump_thread.is_alive(): + self._pump_thread.join(timeout = 2.0) + if self._pump_thread.is_alive(): + self._update_progress(error = "Previous training event pump is still finalizing") + return False + if not self._model_config: + self._update_progress(error = "Model not loaded") + return False + if not self._dataset_config: + self._update_progress(error = "Dataset not loaded") + return False + if self.is_cpt: + self._update_progress( + error = "Continued Pretraining is not supported for MLX training yet.", + is_training = False, + is_completed = False, + ) + return False + + config = self._build_worker_config(training_args) + event_queue = queue.Queue() + stop_queue = queue.Queue() + self._event_queue = event_queue + self._stop_queue = stop_queue + self.should_stop = False + self.is_training = True + self.training_progress = TrainingProgress( + is_training = True, + status_message = "Initializing MLX training...", + ) + + self.training_thread = threading.Thread( + target = self._run_training_thread, + args = (config, event_queue, stop_queue), + daemon = True, + ) + self._pump_thread = threading.Thread( + target = self._pump_events, + args = (event_queue, self.training_thread), + daemon = True, + ) + self.training_thread.start() + self._pump_thread.start() + return True + + def _build_worker_config(self, training_args: dict[str, Any]) -> dict[str, Any]: + peft = { + "use_lora": True, + "lora_r": 16, + "lora_alpha": 16, + "lora_dropout": 0.0, + "target_modules": None, + "gradient_checkpointing": "unsloth", + "use_rslora": False, + "use_loftq": False, + "finetune_vision_layers": True, + "finetune_language_layers": True, + "finetune_attention_modules": True, + "finetune_mlp_modules": True, + **self._peft_config, + } + output_dir = training_args.get("output_dir") + if output_dir: + output_dir = os.path.abspath(os.path.expanduser(str(output_dir))) + values = { + **self._model_config, + **self._dataset_config, + **training_args, + "training_type": ( + "Continued Pretraining" + if self.is_cpt + else "LoRA/QLoRA" + if peft["use_lora"] + else "Full Finetuning" + ), + **peft, + "output_dir": output_dir, + "allow_external_output_dir": bool(output_dir), + } + config = _build_training_worker_config(values) + config["resolved_gpu_ids"] = None + config["gpu_selection"] = None + return config + + def _run_training_thread( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + try: + self._run_mlx_worker(config, event_queue, stop_queue) + except Exception as exc: + if event_queue is not None: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def _run_mlx_worker( + self, config: dict[str, Any], event_queue: queue.Queue, stop_queue: queue.Queue + ): + from .worker import run_mlx_training_process + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + ) + + def _pump_events(self, event_queue: queue.Queue, training_thread: threading.Thread): + while True: + event = None + try: + event = event_queue.get(timeout = 0.25) + except queue.Empty: + pass + if event is not None: + self._handle_event(event) + continue + if not training_thread.is_alive(): + self._drain_events(event_queue) + with self._lock: + if self.training_progress.is_training: + self.training_progress.is_training = False + if self.should_stop: + self.training_progress.status_message = "Training stopped." + elif ( + not self.training_progress.error + and not self.training_progress.is_completed + ): + self.training_progress.error = "Training process exited unexpectedly" + self.is_training = False + self._event_queue = None + self._stop_queue = None + return + + def _drain_events(self, event_queue: Optional[queue.Queue] = None): + event_queue = event_queue or self._event_queue + if event_queue is None: + return + while True: + try: + self._handle_event(event_queue.get_nowait()) + except queue.Empty: + return + + def _handle_event(self, event: dict[str, Any]): + etype = event.get("type") + if etype == "status": + self._update_progress( + status_message = event.get("status_message") or event.get("message") or "" + ) + return + if etype == "progress": + self._update_progress( + step = event.get("step", self.training_progress.step), + epoch = event.get("epoch", self.training_progress.epoch), + loss = event.get("loss", self.training_progress.loss), + learning_rate = event.get("learning_rate", self.training_progress.learning_rate), + total_steps = event.get("total_steps", self.training_progress.total_steps), + elapsed_seconds = event.get( + "elapsed_seconds", + self.training_progress.elapsed_seconds, + ), + eta_seconds = event.get("eta_seconds", self.training_progress.eta_seconds), + grad_norm = event.get("grad_norm", self.training_progress.grad_norm), + num_tokens = event.get("num_tokens", self.training_progress.num_tokens), + eval_loss = event.get("eval_loss", self.training_progress.eval_loss), + peak_memory_gb = event.get("peak_memory_gb", self.training_progress.peak_memory_gb), + ) + return + if etype == "complete": + status_message = event.get("status_message") or "Training completed" + output_dir = event.get("output_dir") + was_cancelled = self.should_stop or status_message.strip().lower() in { + "training cancelled", + "training stopped", + } + self.output_dir = output_dir + self._update_progress( + is_training = False, + is_completed = not was_cancelled, + error = None, + status_message = status_message, + output_dir = output_dir, + ) + self.is_training = False + return + if etype == "error": + self._update_progress( + is_training = False, + is_completed = False, + error = event.get("error") or event.get("message") or "Training failed", + ) + self.is_training = False + return + + def stop_training(self, save: bool = True): + self.should_stop = True + self.save_on_stop = bool(save) + if self._stop_queue is not None: + self._stop_queue.put({"type": "stop", "save": save}) + status_message = ( + "Stopping training and saving checkpoint..." if save else "Cancelling training..." + ) + self._update_progress(status_message = status_message) + return True + + def get_training_progress(self) -> TrainingProgress: + pump_thread = self._pump_thread + training_thread = self.training_thread + if ( + pump_thread is not None + and pump_thread.is_alive() + and (training_thread is None or not training_thread.is_alive()) + and threading.current_thread() is not pump_thread + ): + pump_thread.join(timeout = 5.0) + if pump_thread is None or not pump_thread.is_alive(): + self._drain_events() + with self._lock: + return replace(self.training_progress) + + +def create_mlx_trainer_adapter(*args, **kwargs): + return _MLXTrainerAdapter(*args, **kwargs) class TrainingBackend: @@ -296,86 +816,7 @@ class TrainingBackend: # treat this fresh setup as a recoverable death. self._pump_running = False - # Build config dict for the subprocess - config = { - "model_name": kwargs["model_name"], - "project_name": kwargs.get("project_name"), - "training_type": kwargs.get("training_type", "LoRA/QLoRA"), - "hf_token": kwargs.get("hf_token", ""), - "load_in_4bit": kwargs.get("load_in_4bit", True), - "max_seq_length": kwargs.get("max_seq_length", 2048), - "vision_image_size": kwargs.get("vision_image_size"), - "hf_dataset": kwargs.get("hf_dataset", ""), - "local_datasets": kwargs.get("local_datasets"), - "local_eval_datasets": kwargs.get("local_eval_datasets"), - "format_type": kwargs.get("format_type", ""), - "subset": kwargs.get("subset"), - "train_split": kwargs.get("train_split", "train"), - "eval_split": kwargs.get("eval_split"), - "eval_steps": kwargs.get("eval_steps", 0.00), - "dataset_streaming": kwargs.get("dataset_streaming", False), - "dataset_slice_start": kwargs.get("dataset_slice_start"), - "dataset_slice_end": kwargs.get("dataset_slice_end"), - "custom_format_mapping": kwargs.get("custom_format_mapping"), - "is_dataset_image": kwargs.get("is_dataset_image", False), - "is_dataset_audio": kwargs.get("is_dataset_audio", False), - "is_embedding": kwargs.get("is_embedding", False), - "num_epochs": kwargs.get("num_epochs", 3), - "learning_rate": kwargs.get("learning_rate", "2e-4"), - "embedding_learning_rate": kwargs.get("embedding_learning_rate"), - "batch_size": kwargs.get("batch_size", 2), - "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), - "warmup_steps": kwargs.get("warmup_steps"), - "warmup_ratio": kwargs.get("warmup_ratio"), - "max_steps": kwargs.get("max_steps", 0), - "save_steps": kwargs.get("save_steps", 0), - "weight_decay": kwargs.get("weight_decay", 0.001), - "max_grad_norm": kwargs.get("max_grad_norm", 0.0), - "max_grad_value": _coerce_optional_nonneg_float( - "max_grad_value", kwargs.get("max_grad_value") - ), - "max_grad_leaf_norm": _coerce_optional_nonneg_float( - "max_grad_leaf_norm", kwargs.get("max_grad_leaf_norm") - ), - "cast_norm_output_to_input_dtype": _coerce_optional_bool( - kwargs.get("cast_norm_output_to_input_dtype"), True - ), - # MLX/CUDA/embedding workers need an int (transformers.set_seed(None) raises). - "random_seed": _coerce_seed(kwargs.get("random_seed")), - "packing": kwargs.get("packing", False), - "optim": kwargs.get("optim", "adamw_8bit"), - "lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"), - "use_lora": kwargs.get("use_lora", True), - "lora_r": kwargs.get("lora_r", 16), - "lora_alpha": kwargs.get("lora_alpha", 16), - "lora_dropout": kwargs.get("lora_dropout", 0.0), - "target_modules": kwargs.get("target_modules"), - "gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"), - "use_rslora": kwargs.get("use_rslora", False), - "use_loftq": kwargs.get("use_loftq", False), - "train_on_completions": kwargs.get("train_on_completions", False), - "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), - "finetune_language_layers": kwargs.get("finetune_language_layers", True), - "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), - "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), - "enable_wandb": kwargs.get("enable_wandb", False), - "wandb_token": kwargs.get("wandb_token"), - "wandb_project": kwargs.get("wandb_project", "unsloth-training"), - "enable_tensorboard": kwargs.get("enable_tensorboard", False), - "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), - "resume_from_checkpoint": kwargs.get("resume_from_checkpoint"), - "trust_remote_code": kwargs.get("trust_remote_code", False), - "approved_remote_code_fingerprint": kwargs.get("approved_remote_code_fingerprint"), - "subject": kwargs.get("subject"), - "gpu_ids": kwargs.get("gpu_ids"), - "s3_config": kwargs.get("s3_config"), - # Flipped to True only by the HTTP-fallback respawn after a stall. - "disable_xet": kwargs.get("disable_xet", False), - } - - # Full finetuning always runs in 16-bit; LoRA/QLoRA/CPT keep the request. - if config["training_type"] == "Full Finetuning": - config["load_in_4bit"] = False + config = _build_training_worker_config(kwargs) # Split GPU validation from placement around the VRAM hook: # * Explicit gpu_ids are validated here (raises -> the route returns 400 @@ -401,7 +842,7 @@ class TrainingBackend: ) defer_auto_selection = False - if _hw.DEVICE == _hw.DeviceType.MLX: + if should_use_mlx_training_backend(device = _hw.DEVICE): config["resolved_gpu_ids"] = None config["gpu_selection"] = None elif gpu_ids: @@ -1022,17 +1463,22 @@ class TrainingBackend: self._progress.is_training = True elif etype == "complete": - self._progress.is_training = False - self._progress.is_completed = True - self._output_dir = event.get("output_dir") msg = event.get("status_message", "Training completed") + stopped = self._should_stop or msg.strip().lower() in { + "training cancelled", + "training stopped", + } + self._progress.is_training = False + self._progress.is_completed = not stopped + self._output_dir = event.get("output_dir") + self._progress.output_dir = self._output_dir self._progress.status_message = msg if not self._db_run_created and self.current_job_id and self._db_config: db_action = "create_and_finalize" else: db_action = "finalize" db_action_kwargs = { - "status": "stopped" if self._should_stop else "completed", + "status": "stopped" if stopped else "completed", "output_dir": self._output_dir, } diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 17dc1299ca..0ff4d517ed 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -1309,14 +1309,18 @@ def _normalize_mlx_studio_scheduler(value): def _resolve_mlx_local_dataset_files(file_paths: list) -> list[str]: - """Resolve Studio local dataset uploads without importing the GPU trainer.""" + """Resolve CLI paths and Studio local dataset uploads without importing the GPU trainer.""" from utils.paths import resolve_dataset_path all_files: list[str] = [] for dataset_file in file_paths or []: - file_path = ( - dataset_file if os.path.isabs(dataset_file) else str(resolve_dataset_path(dataset_file)) - ) + dataset_path = Path(os.path.expanduser(str(dataset_file))) + if dataset_path.is_absolute(): + file_path = str(dataset_path) + elif dataset_path.exists(): + file_path = str(dataset_path.resolve()) + else: + file_path = str(resolve_dataset_path(str(dataset_file))) file_path_obj = Path(file_path) if file_path_obj.is_dir(): @@ -1355,6 +1359,58 @@ def _mlx_local_dataset_loader_for_files(files: list[str]) -> str: raise ValueError(f"Unsupported dataset format: {files[0]}") +_MLX_WORKER_COMPLETE = "_mlx_worker_complete" + + +def _start_mlx_stop_poller(stop_queue): + import queue as _queue + import threading + + stop_save = [True] + stop_requested = [False] + trainer_ref = [None] + + def is_stop_requested(): + return stop_requested[0] + + def poll_stop(): + while True: + try: + msg = stop_queue.get(timeout = 0.25) + if msg and msg.get("type") == _MLX_WORKER_COMPLETE: + return + if msg and msg.get("type") == "stop": + stop_save[0] = msg.get("save", True) + stop_requested[0] = True + trainer = trainer_ref[0] + if trainer is not None: + trainer.stop_requested = True + return + except _queue.Empty: + continue + except (EOFError, OSError): + return + + stop_thread = threading.Thread(target = poll_stop, daemon = True) + stop_thread.start() + return stop_save, stop_requested, trainer_ref, is_stop_requested, stop_thread + + +def _resolve_mlx_output_dir(config, model_name): + from utils.paths import resolve_output_dir, default_run_dir_name + + output_dir = config.get("output_dir", "") + if not output_dir: + output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}" + return str(resolve_output_dir(output_dir)) + if config.get("allow_external_output_dir"): + output_path = Path(output_dir).expanduser() + if not output_path.is_absolute(): + output_path = Path.cwd() / output_path + return str(output_path.resolve()) + return str(resolve_output_dir(output_dir)) + + def _run_mlx_training(event_queue, stop_queue, config): """Self-contained MLX training path for Apple Silicon. @@ -1363,8 +1419,6 @@ def _run_mlx_training(event_queue, stop_queue, config): """ import time import math - import threading - import queue as _queue from pathlib import Path def _send(event_type, **kwargs): @@ -1374,31 +1428,9 @@ def _run_mlx_training(event_queue, stop_queue, config): kwargs["message"] = sm event_queue.put({"type": event_type, "ts": time.time(), **kwargs}) - _stop_save = [True] - _stop_requested = [False] - _trainer_ref = [None] - - def _is_stop_requested(): - return _stop_requested[0] - - def _poll_stop(): - while True: - try: - msg = stop_queue.get(timeout = 1.0) - if msg and msg.get("type") == "stop": - _stop_save[0] = msg.get("save", True) - _stop_requested[0] = True - trainer = _trainer_ref[0] - if trainer is not None: - trainer.stop_requested = True - return - except _queue.Empty: - continue - except (EOFError, OSError): - return - - stop_thread = threading.Thread(target = _poll_stop, daemon = True) - stop_thread.start() + _stop_save, _stop_requested, _trainer_ref, _is_stop_requested, _stop_thread = ( + _start_mlx_stop_poller(stop_queue) + ) _send("status", status_message = "Loading MLX libraries...") @@ -1804,21 +1836,14 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 5. Build output dir ── # Resolve to ~/.unsloth/studio/outputs/ so the export page finds it - from utils.paths import resolve_output_dir, ensure_dir + from utils.paths import ensure_dir - output_dir = config.get("output_dir", "") - if not output_dir: - output_dir = build_default_output_dir_name( - model_name, - config.get("project_name"), - ) - output_dir = str(resolve_output_dir(output_dir)) + output_dir = _resolve_mlx_output_dir(config, model_name) ensure_dir(Path(output_dir)) # ── 6. Create trainer ── eval_steps_val = config.get("eval_steps", 0) or 0 if isinstance(eval_steps_val, float) and 0 < eval_steps_val < 1: - # Studio sometimes sends fraction-of-total-steps eval_steps_val = max(1, int(eval_steps_val * max_steps)) else: eval_steps_val = int(eval_steps_val) @@ -2043,12 +2068,27 @@ def _run_mlx_training(event_queue, stop_queue, config): # ── 11. Run training ── gc.collect() mx.synchronize() - trainer.train(resume_from_checkpoint = resume_from_checkpoint) + _save_model = trainer.save_model + + def _skip_internal_final_save(*args, **kwargs): + raise ValueError("worker owns final save") + + trainer.save_model = _skip_internal_final_save + try: + trainer.train(resume_from_checkpoint = resume_from_checkpoint) + finally: + trainer.save_model = _save_model # ── 12. Save and finalize ── - if trainer.stop_requested and not _stop_save[0]: - # User clicked "Cancel" (save=False) — skip saving - _send("complete", output_dir = None, status_message = "Training cancelled") + if trainer.stop_requested: + if not _stop_save[0]: + # Cancel (save=False): skip saving. + _send("complete", output_dir = None, status_message = "Training cancelled") + else: + _send("status", status_message = "Saving stopped model...") + mx.synchronize() + trainer.save_model(output_dir) + _send("complete", output_dir = output_dir, status_message = "Training stopped") else: _send("status", status_message = "Saving model...") mx.synchronize() @@ -2067,6 +2107,79 @@ def _run_mlx_training(event_queue, stop_queue, config): pass +def _is_current_process_apple_silicon() -> bool: + import platform + return platform.system() == "Darwin" and platform.machine() == "arm64" + + +def run_mlx_training_process( + *, + event_queue: Any, + stop_queue: Any, + config: dict, + transformers_activated: bool = False, +) -> None: + """MLX worker entrypoint shared by Studio subprocesses and the CLI adapter.""" + model_name = config["model_name"] + + backend_path = str(Path(__file__).resolve().parent.parent.parent) + if backend_path not in sys.path: + sys.path.insert(0, backend_path) + + from utils.hf_xet_fallback import child_should_disable_xet + + if child_should_disable_xet(config): + os.environ["HF_HUB_DISABLE_XET"] = "1" + os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "0" + + if not transformers_activated: + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE != _hw.DeviceType.MLX: + event_queue.put( + { + "type": "error", + "error": "MLX training requires Apple Silicon with the MLX backend available.", + "stack": "", + "ts": time.time(), + } + ) + return + + if config.get("is_dataset_audio"): + event_queue.put( + { + "type": "error", + "error": "Audio dataset training is not yet supported on Apple Silicon.", + "stack": "", + "ts": time.time(), + } + ) + return + + try: + try: + _run_mlx_training(event_queue, stop_queue, config) + finally: + try: + stop_queue.put({"type": _MLX_WORKER_COMPLETE}) + except (EOFError, OSError, ValueError): + pass + except Exception as exc: + event_queue.put( + { + "type": "error", + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + "ts": time.time(), + } + ) + + def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> None: """Subprocess entrypoint. Fresh Python — no stale module state. @@ -2141,36 +2254,26 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) + from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend + + mlx_backend_requested = is_apple_silicon_training_platform() + + mlx_transformers_activated = False + if mlx_backend_requested and _is_current_process_apple_silicon(): + # Must precede detect_hardware(): its MLX stack check imports mlx_lm, hence transformers. + _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) + mlx_transformers_activated = True + from utils.hardware import hardware as _hw _hw.detect_hardware() - if _hw.DEVICE == _hw.DeviceType.MLX: - if config.get("is_dataset_audio"): - event_queue.put( - { - "type": "error", - "error": "Audio dataset training is not yet supported on Apple Silicon.", - "stack": "", - "ts": time.time(), - } - ) - return - # Activate correct transformers version (Gemma-4 needs a 5.x sidecar, etc.) - # Must happen before any transformers/mlx-lm imports in _run_mlx_training. - # Non-fatal: fall through with whatever version is installed, but log - # the failure instead of swallowing it (issue #6103). - _activate_transformers_version_or_warn(model_name, config.get("hf_token") or None) - try: - _run_mlx_training(event_queue, stop_queue, config) - except Exception as exc: - event_queue.put( - { - "type": "error", - "error": str(exc), - "stack": traceback.format_exc(limit = 20), - "ts": time.time(), - } - ) + if mlx_backend_requested or should_use_mlx_training_backend(device = _hw.DEVICE): + run_mlx_training_process( + event_queue = event_queue, + stop_queue = stop_queue, + config = config, + transformers_activated = mlx_transformers_activated, + ) return # ── 1. Activate correct transformers version BEFORE any ML imports ── @@ -2693,7 +2796,8 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> if backend_path not in sys.path: sys.path.insert(0, backend_path) - from core.training.trainer import UnslothTrainer, TrainingProgress + from core.training.training import TrainingProgress + from core.training.trainer import UnslothTrainer from utils.paths import ( ensure_dir, resolve_output_dir, diff --git a/studio/backend/tests/test_training_preflight.py b/studio/backend/tests/test_training_preflight.py index 54048a65dd..47c6669f8f 100644 --- a/studio/backend/tests/test_training_preflight.py +++ b/studio/backend/tests/test_training_preflight.py @@ -6,9 +6,15 @@ empty-chat-template crash) before train(). The real methods are bound onto a lig fake self so the production logic runs against controlled batches.""" import importlib +import json +import os +import queue +import subprocess import sys +import threading import types import unittest +from pathlib import Path from types import SimpleNamespace from unittest.mock import MagicMock @@ -184,5 +190,231 @@ class TestChatTemplateRendersEmpty(unittest.TestCase): self.assertFalse(s._chat_template_renders_empty()) +def _clear_trainer_module(package: str): + sys.modules.pop(f"{package}.trainer", None) + pkg = sys.modules.get(package) + if pkg is not None and hasattr(pkg, "trainer"): + delattr(pkg, "trainer") + + +def _set_training_platform(monkeypatch, package: str, backend: str): + training_mod = importlib.import_module(f"{package}.training") + from utils.hardware import hardware as hw + + monkeypatch.setattr(hw, "DEVICE", None) + monkeypatch.setattr( + training_mod.platform, + "system", + lambda: "Darwin" if backend == "mlx" else "Linux", + ) + monkeypatch.setattr( + training_mod.platform, + "machine", + lambda: "arm64" if backend == "mlx" else "x86_64", + ) + + +def _load_trainer_module( + monkeypatch, + backend: str, + package: str = "core.training", +): + _set_training_platform(monkeypatch, package, backend) + _clear_trainer_module(package) + if package in sys.modules: + importlib.reload(sys.modules[package]) + trainer_mod = importlib.import_module(f"{package}.trainer") + training_mod = importlib.import_module(f"{package}.training") + monkeypatch.setattr( + training_mod._MLXTrainerAdapter, + "_activate_transformers_for_model", + lambda self, model_name, hf_token: None, + ) + return trainer_mod + + +class _ExitedProc: + def join(self, timeout = None): + return None + + def is_alive(self): + return False + + +class _TerminableProc: + def __init__(self): + self.terminated = False + self._done = threading.Event() + + def join(self, timeout = None): + self._done.wait(timeout = timeout or 5) + + def is_alive(self): + return not self.terminated + + def terminate(self): + self.terminated = True + self._done.set() + + +def test_unsloth_trainer_dispatches_for_mlx_and_torch(monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + + mlx_trainer = trainer_mod.UnslothTrainer() + + assert type(mlx_trainer).__module__ == "core.training.training" + assert mlx_trainer.get_training_progress().status_message == "Ready to train" + + trainer_mod = _load_trainer_module(monkeypatch, "torch") + + assert trainer_mod.UnslothTrainer().__class__ is trainer_mod.UnslothTrainer + + +def test_cli_mlx_trainer_activates_before_importing_trainer(): + repo_root = Path(__file__).resolve().parents[3] + script = """ +import json +import sys +import unsloth_cli.commands.train as train_cmd +from studio.backend.core.training import training as training_mod +from utils.hardware import hardware as hw + +training_mod.platform.system = lambda: "Darwin" +training_mod.platform.machine = lambda: "arm64" +hw.DEVICE = None +events = [] + +def fake_activate(model_name, hf_token): + events.append({ + "model_name": model_name, + "trainer_loaded": "studio.backend.core.training.trainer" in sys.modules, + }) + +train_cmd._activate_mlx_transformers = fake_activate +trainer = train_cmd._create_cli_trainer("mlx-community/Qwen3-0.6B-4bit", None) +print(json.dumps({ + "trainer_module": type(trainer).__module__, + "events": events, +})) +""" + env = os.environ.copy() + env["PYTHONPATH"] = os.pathsep.join( + [str(repo_root), str(repo_root / "studio" / "backend"), env.get("PYTHONPATH", "")] + ) + result = subprocess.run( + [sys.executable, "-c", script], + cwd = repo_root, + env = env, + text = True, + stdout = subprocess.PIPE, + stderr = subprocess.PIPE, + check = True, + ) + payload = json.loads(result.stdout) + + assert payload["trainer_module"] == "studio.backend.core.training.training" + assert payload["events"] == [ + {"model_name": "mlx-community/Qwen3-0.6B-4bit", "trainer_loaded": False} + ] + + +def test_mlx_adapter_builds_config_and_reports_completion(tmp_path, monkeypatch): + trainer_mod = _load_trainer_module(monkeypatch, "mlx") + captured = {} + + def fake_run_worker(config, event_queue, stop_queue): + captured["config"] = config + event_queue.put({"type": "progress", "step": 1, "total_steps": 1, "loss": 0.25}) + event_queue.put( + {"type": "complete", "status_message": "done", "output_dir": config["output_dir"]} + ) + + trainer = trainer_mod.UnslothTrainer() + monkeypatch.setattr(trainer, "_run_mlx_worker", fake_run_worker) + + assert trainer.load_model("mlx-community/Qwen3-0.6B-4bit", max_seq_length = 1024) + assert trainer.prepare_model_for_training(use_lora = False) + dataset, eval_dataset = trainer.load_and_format_dataset("org/dataset") + output_dir = tmp_path / "mlx-out" + + assert trainer.start_training( + dataset = dataset, + eval_dataset = eval_dataset, + output_dir = output_dir, + project_name = "Sales Assistant", + max_steps = 1, + learning_rate = 3e-4, + ) + trainer.training_thread.join(timeout = 5) + + progress = trainer.get_training_progress() + config = captured["config"] + assert progress.is_completed + assert progress.output_dir == str(output_dir.resolve()) + progress.status_message = "mutated" + assert trainer.get_training_progress().status_message == "done" + assert config["model_name"] == "mlx-community/Qwen3-0.6B-4bit" + assert config["project_name"] == "Sales Assistant" + assert config["hf_dataset"] == "org/dataset" + assert config["training_type"] == "Full Finetuning" + assert config["load_in_4bit"] is False + assert config["max_seq_length"] == 1024 + assert config["learning_rate"] == 3e-4 + assert config["output_dir"] == str(output_dir.resolve()) + assert config["allow_external_output_dir"] is True + + +def test_mlx_worker_helpers_cover_cli_paths(tmp_path, monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training.worker import ( + _resolve_mlx_local_dataset_files, + _resolve_mlx_output_dir, + ) + + dataset = tmp_path / "train.jsonl" + dataset.write_text('{"text":"hello"}\n', encoding = "utf-8") + monkeypatch.chdir(tmp_path) + + assert _resolve_mlx_local_dataset_files(["train.jsonl"]) == [str(dataset)] + assert _resolve_mlx_output_dir( + {"output_dir": "cli-out", "allow_external_output_dir": True}, + "mlx-community/Qwen3-0.6B-4bit", + ) == str((tmp_path / "cli-out").resolve()) + + +def test_run_mlx_training_process_applies_side_effects_before_hardware_detection(monkeypatch): + _load_trainer_module(monkeypatch, "mlx") + from core.training import worker + from utils.hardware import hardware as hw + + order = [] + + def fake_activate(model_name, hf_token): + order.append(("activate", model_name, hf_token)) + + def fake_detect_hardware(): + order.append("detect") + hw.DEVICE = hw.DeviceType.CPU + return hw.DEVICE + + monkeypatch.delenv("HF_HUB_DISABLE_XET", raising = False) + monkeypatch.delenv("HF_HUB_ENABLE_HF_TRANSFER", raising = False) + monkeypatch.setattr(worker, "_activate_transformers_version_or_warn", fake_activate) + monkeypatch.setattr(hw, "detect_hardware", fake_detect_hardware) + + event_queue = queue.Queue() + worker.run_mlx_training_process( + event_queue = event_queue, + stop_queue = queue.Queue(), + config = {"model_name": "mlx-community/Gemma-4-12B", "disable_xet": True}, + ) + + event = event_queue.get_nowait() + assert order == [("activate", "mlx-community/Gemma-4-12B", None), "detect"] + assert os.environ["HF_HUB_DISABLE_XET"] == "1" + assert os.environ["HF_HUB_ENABLE_HF_TRANSFER"] == "0" + assert "MLX training requires Apple Silicon" in event["error"] + + if __name__ == "__main__": unittest.main() diff --git a/unsloth-cli.py b/unsloth-cli.py index 756efef0d0..be48893d8d 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -22,24 +22,179 @@ import argparse import os +def _is_mlx_backend(unsloth_module): + return bool(getattr(unsloth_module, "_IS_MLX", False)) + + +def _normalize_dtype(dtype, is_mlx): + if is_mlx and isinstance(dtype, str) and dtype.strip().lower() in {"", "none", "auto"}: + return None + return dtype + + +def _prepare_device_map(is_mlx): + if is_mlx: + return None, False + + from unsloth.models.loader_utils import prepare_device_map + return prepare_device_map() + + +class _CallableTokenizerProxy: + def __init__(self, tokenizer): + self._tokenizer = tokenizer + + def __getattr__(self, name): + return getattr(self._tokenizer, name) + + def __call__(self, text, *args, **kwargs): + # MLX/torch-free: never request torch tensors; keep plain python ids. + kwargs.pop("return_tensors", None) + wrapped = getattr(self._tokenizer, "_tokenizer", None) + if callable(wrapped): + return wrapped(text, *args, **kwargs) + + add_special_tokens = kwargs.get("add_special_tokens", False) + input_ids = self._tokenizer.encode(text, add_special_tokens = add_special_tokens) + return {"input_ids": input_ids} + + +def _tokenizer_for_raw_text_loader(tokenizer, is_mlx): + if not is_mlx or callable(tokenizer): + return tokenizer + return _CallableTokenizerProxy(tokenizer) + + +def _raw_text_loader_for_backend( + RawTextDataLoader, + tokenizer, + is_mlx, + chunk_size = 2048, + stride = 512, +): + return RawTextDataLoader( + _tokenizer_for_raw_text_loader(tokenizer, is_mlx), + chunk_size, + stride, + return_tokenized = not is_mlx, + ) + + +def _train_with_legacy_save_control(trainer, is_mlx): + if not is_mlx: + return trainer.train() + + original_save_model = getattr(trainer, "save_model", None) + if original_save_model is None: + return trainer.train() + + def skip_internal_final_save(*args, **kwargs): + raise ValueError("legacy unsloth-cli.py owns final save") + + trainer.save_model = skip_internal_final_save + try: + return trainer.train() + finally: + trainer.save_model = original_save_model + + +def _iter_quantization_methods(quantization): + if isinstance(quantization, list): + return quantization + return [quantization] + + +def _save_or_push_model(model, tokenizer, args, is_mlx): + if not args.save_model: + print("Warning: The model is not saved!") + return + + # Enter the GGUF branch when saving or pushing GGUF, so --push_gguf works + # without --save_gguf (the local save is guarded separately below). + if args.save_gguf or args.push_gguf: + if not args.save_gguf: + print("Warning: --save_gguf not set, pushing GGUF to hub without saving locally.") + for quantization_method in _iter_quantization_methods(args.quantization): + if args.save_gguf: + print(f"Saving model with quantization method: {quantization_method}") + model.save_pretrained_gguf( + args.save_path, + tokenizer, + quantization_method = quantization_method, + ) + if args.push_model or args.push_gguf: + model.push_to_hub_gguf( + args.hub_path, + tokenizer, + quantization_method = quantization_method, + token = args.hub_token, + ) + return + + if is_mlx: + model.save_pretrained_merged( + args.save_path, + tokenizer, + save_method = args.save_method, + push_to_hub = args.push_model, + repo_id = args.hub_path if args.push_model else None, + token = args.hub_token, + ) + return + + model.save_pretrained_merged(args.save_path, tokenizer, save_method = args.save_method) + if args.push_model: + model.push_to_hub_merged(args.hub_path, tokenizer, args.save_method, token = args.hub_token) + + +def _build_sft_config(SFTConfig, args, is_mlx, bf16_supported): + config_kwargs = dict( + per_device_train_batch_size = args.per_device_train_batch_size, + gradient_accumulation_steps = args.gradient_accumulation_steps, + warmup_steps = args.warmup_steps, + max_steps = args.max_steps, + learning_rate = args.learning_rate, + fp16 = not bf16_supported, + bf16 = bf16_supported, + logging_steps = args.logging_steps, + optim = args.optim, + weight_decay = args.weight_decay, + lr_scheduler_type = args.lr_scheduler_type, + seed = args.seed, + output_dir = args.output_dir, + report_to = args.report_to, + max_length = args.max_seq_length, + dataset_num_proc = 2, + packing = args.packing, + ) + if is_mlx: + if args.per_device_eval_batch_size != 4: + print("Warning: --per_device_eval_batch_size is ignored on MLX without eval data.") + else: + config_kwargs["per_device_eval_batch_size"] = args.per_device_eval_batch_size + return SFTConfig(**config_kwargs) + + def run(args): + import unsloth from unsloth import FastLanguageModel from datasets import load_dataset from transformers.utils import strtobool from trl import SFTTrainer, SFTConfig from unsloth import is_bfloat16_supported - from unsloth.models.loader_utils import prepare_device_map import logging from unsloth import RawTextDataLoader logging.getLogger("hf-to-gguf").setLevel(logging.WARNING) + is_mlx = _is_mlx_backend(unsloth) + # Load model and tokenizer - device_map, distributed = prepare_device_map() + device_map, distributed = _prepare_device_map(is_mlx) model, tokenizer = FastLanguageModel.from_pretrained( model_name = args.model_name, max_seq_length = args.max_seq_length, - dtype = args.dtype, + dtype = _normalize_dtype(args.dtype, is_mlx), load_in_4bit = args.load_in_4bit, device_map = device_map, ) @@ -92,11 +247,13 @@ def run(args): def load_dataset_smart(args): from transformers.utils import strtobool if args.raw_text_file: - loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride) + loader = _raw_text_loader_for_backend( + RawTextDataLoader, tokenizer, is_mlx, args.chunk_size, args.stride + ) dataset = loader.load_from_file(args.raw_text_file) elif args.dataset.endswith((".txt", ".md", ".json", ".jsonl")): # Auto-detect local raw text files - loader = RawTextDataLoader(tokenizer) + loader = _raw_text_loader_for_backend(RawTextDataLoader, tokenizer, is_mlx) dataset = loader.load_from_file(args.dataset) else: use_modelscope = strtobool(os.environ.get("UNSLOTH_USE_MODELSCOPE", "False")) @@ -115,27 +272,9 @@ def run(args): print("Data is formatted and ready!") # Configure training arguments - training_args = SFTConfig( - per_device_train_batch_size = args.per_device_train_batch_size, - per_device_eval_batch_size = args.per_device_eval_batch_size, - gradient_accumulation_steps = args.gradient_accumulation_steps, - warmup_steps = args.warmup_steps, - max_steps = args.max_steps, - learning_rate = args.learning_rate, - fp16 = not is_bfloat16_supported(), - bf16 = is_bfloat16_supported(), - logging_steps = args.logging_steps, - optim = args.optim, - weight_decay = args.weight_decay, - lr_scheduler_type = args.lr_scheduler_type, - seed = args.seed, - output_dir = args.output_dir, - report_to = args.report_to, - max_length = args.max_seq_length, - dataset_num_proc = 2, - ddp_find_unused_parameters = False if distributed else None, - packing = args.packing, - ) + training_args = _build_sft_config(SFTConfig, args, is_mlx, is_bfloat16_supported()) + if distributed: + training_args.ddp_find_unused_parameters = False # Initialize trainer trainer = SFTTrainer( @@ -145,57 +284,9 @@ def run(args): args = training_args, ) - trainer.train() + _train_with_legacy_save_control(trainer, is_mlx) - # Save model - if args.save_model: - # If args.quantization is a list, save once per quantization method - # Enter the GGUF branch when saving *or* pushing GGUF, so --push_gguf - # works even when --save_gguf is omitted (the local save is guarded - # separately below). - if args.save_gguf or args.push_gguf: - # Push-only GGUF (no --save_gguf) skips the local save; warn so it is not silent. - if not args.save_gguf: - print("Warning: --save_gguf not set, pushing GGUF to hub without saving locally.") - if isinstance(args.quantization, list): - for quantization_method in args.quantization: - if args.save_gguf: - print(f"Saving model with quantization method: {quantization_method}") - model.save_pretrained_gguf( - args.save_path, - tokenizer, - quantization_method = quantization_method, - ) - if args.push_model or args.push_gguf: - model.push_to_hub_gguf( - args.hub_path, - tokenizer, - quantization_method = quantization_method, - token = args.hub_token, - ) - else: - if args.save_gguf: - print(f"Saving model with quantization method: {args.quantization}") - model.save_pretrained_gguf( - args.save_path, - tokenizer, - quantization_method = args.quantization, - ) - if args.push_model or args.push_gguf: - model.push_to_hub_gguf( - args.hub_path, - tokenizer, - quantization_method = args.quantization, - token = args.hub_token, - ) - else: - model.save_pretrained_merged(args.save_path, tokenizer, args.save_method) - if args.push_model: - model.push_to_hub_merged( - args.hub_path, tokenizer, args.save_method, token = args.hub_token - ) - else: - print("Warning: The model is not saved!") + _save_or_push_model(model, tokenizer, args, is_mlx) if __name__ == "__main__": diff --git a/unsloth_cli/commands/train.py b/unsloth_cli/commands/train.py index 9d47a574d5..c52c2344b7 100644 --- a/unsloth_cli/commands/train.py +++ b/unsloth_cli/commands/train.py @@ -7,10 +7,42 @@ from typing import Optional import typer +from unsloth_cli._inference import ensure_studio_backend_path from unsloth_cli.config import Config, load_config from unsloth_cli.options import add_options_from_config +def _should_use_mlx_backend_for_cli() -> bool: + ensure_studio_backend_path() + from studio.backend.core.training.training import should_use_mlx_training_backend + return should_use_mlx_training_backend() + + +def _activate_mlx_transformers(model_name: str, hf_token: Optional[str]) -> None: + # Activate before any transformers import: adapter model-type detection imports utils.models. + ensure_studio_backend_path() + from utils.transformers_version import activate_transformers_for_subprocess + try: + activate_transformers_for_subprocess(model_name, hf_token) + except Exception as exc: + typer.echo(f"Warning: failed to activate Transformers sidecar: {exc}", err = True) + + +def _create_cli_trainer(model_name: str, hf_token: Optional[str]): + if _should_use_mlx_backend_for_cli(): + _activate_mlx_transformers(model_name, hf_token) + # MLX is torch-free: use the lightweight adapter, not trainer.py (imports torch/unsloth/trl at load). + ensure_studio_backend_path() + from studio.backend.core.training.training import create_mlx_trainer_adapter + + return create_mlx_trainer_adapter() + + ensure_studio_backend_path() + from studio.backend.core.training.trainer import UnslothTrainer + + return UnslothTrainer() + + @add_options_from_config(Config) def train( config: Optional[Path] = typer.Option( @@ -39,6 +71,7 @@ def train( typer.echo(f"Error: {e}", err = True) raise typer.Exit(code = 2) + config_overrides = config_overrides or {} cfg.apply_overrides(**config_overrides) # CLI/env tokens take precedence; guard against unresolved typer.Option @@ -83,9 +116,7 @@ def train( ) raise typer.Exit(code = 2) - from studio.backend.core.training.trainer import UnslothTrainer - - trainer = UnslothTrainer() + trainer = _create_cli_trainer(cfg.model, hf_token) # Load model (trainer.is_vlm is set after this) if not trainer.load_model( @@ -124,13 +155,20 @@ def train( try: while trainer.training_thread and trainer.training_thread.is_alive(): + progress = trainer.get_training_progress() + if getattr(progress, "error", None): + break time.sleep(1) except KeyboardInterrupt: typer.echo("Stopping training (Ctrl+C detected)...") trainer.stop_training() finally: if trainer.training_thread: - trainer.training_thread.join() + progress = trainer.get_training_progress() + if getattr(progress, "error", None): + trainer.training_thread.join(timeout = 5) + else: + trainer.training_thread.join() final = trainer.get_training_progress() if getattr(final, "error", None): From 2a6abe2ff5c643ee853f2e0ef632b1e80d06918a Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:39 +0800 Subject: [PATCH 017/402] feat(cli): support MLX distributed inference (#6845) * feat(cli): detect MLX distributed launch context * feat(mlx): wire distributed inference backend * feat(cli): broadcast MLX distributed chat turns * fix(cli): wait indefinitely for distributed chat turns * fix(cli): report MLX distributed load errors cleanly * fix(mlx): route distributed vlm through loader * fix(cli): detect inline MLX host JSON * fix(studio): harden distributed object sharing * fix(studio): select JACCL distributed backend * fix(cli): abort distributed error paths * Distinguish real stream errors from model text via GenStreamError in distributed CLI * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fail loud when MLX distributed init returns a singleton group The worker only reaches this block when distributed was explicitly requested. A singleton (size 1) group means the launch failed to form a real group (MLX built without distributed support, or an invalid launch env/hostfile); silently continuing leaves nonzero ranks looping forever on share_distributed_object. Raise instead so the surrounding handler returns a clear load error. * Tighten MLX distributed inference comments --------- Co-authored-by: Daniel Han Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../backend/core/inference/mlx_inference.py | 171 ++++++--- studio/backend/core/inference/orchestrator.py | 106 +++++- studio/backend/core/inference/worker.py | 123 ++++++- studio/backend/routes/inference.py | 46 +++ .../tests/test_mlx_inference_backend.py | 125 +++++++ unsloth_cli/_inference.py | 217 ++++++++++-- unsloth_cli/commands/chat.py | 178 +++++++--- unsloth_cli/commands/inference.py | 47 ++- unsloth_cli/tests/test_inference_chat.py | 330 +++++++++++++++++- 9 files changed, 1188 insertions(+), 155 deletions(-) diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index 62d268e15f..d84baa278d 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -5,6 +5,7 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm instead of torch/transformers for model loading and generation. """ +import os import threading from typing import Optional, Generator from core.inference.runtime_context import runtime_context_length @@ -41,6 +42,48 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps): } +def _mlx_distributed_rank_size(group = None): + """Return ``(rank, world_size)`` for an optional MLX distributed group.""" + if group is None: + return 0, 1 + rank = int(group.rank()) + world_size = int(group.size()) + if world_size < 1: + raise ValueError(f"Invalid MLX distributed world_size={world_size}.") + if rank < 0 or rank >= world_size: + raise ValueError(f"Invalid MLX distributed rank={rank} for world_size={world_size}.") + return rank, world_size + + +def _mlx_distributed_backend_from_env(): + if os.environ.get("MLX_JACCL_COORDINATOR") and os.environ.get("MLX_IBV_DEVICES"): + return "jaccl" + return None + + +def _init_mlx_distributed(): + """Initialize MLX distributed state, falling back to singleton metadata.""" + import mlx.core as mx + + group = None + rank = 0 + world_size = 1 + distributed = getattr(mx, "distributed", None) + init = getattr(distributed, "init", None) if distributed is not None else None + if callable(init): + backend = _mlx_distributed_backend_from_env() + if backend is None: + group = init() + else: + try: + group = init(backend = backend) + except TypeError: + group = init() + if group is not None: + rank, world_size = _mlx_distributed_rank_size(group) + return group, rank, world_size + + def _make_mlx_presence_penalty_processor(penalty: float): """Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path. @@ -52,7 +95,7 @@ def _make_mlx_presence_penalty_processor(penalty: float): def _processor(tokens, logits): if state["prompt_len"] is None: - # First call = prompt only; latch its length. + # First call is prompt-only; latch its length. state["prompt_len"] = int(tokens.shape[0]) return logits generated = tokens[state["prompt_len"] :] @@ -61,22 +104,17 @@ def _make_mlx_presence_penalty_processor(penalty: float): import mlx.core as mx vocab = logits.shape[-1] - # Bound generated ids to the valid range [0, vocab) before they index - # logits. MLX does no bounds checking and out-of-bounds indexing is - # documented undefined behavior (crash / memory corruption), unlike the - # torch path's harmless negative wrap -- so this bound is load-bearing - # here and matches the torch filter seen[(seen >= 0) & (seen < vocab)]. - # MLX has no boolean-mask filtering (data-dependent output shape is - # unsupported), so instead of compacting the id list we route every - # out-of-range or negative id to a scratch slot at index ``vocab`` that - # is dropped before the subtract. That scratch slot can never collide - # with a real token, so real ids (including id 0) are penalized exactly - # once and stray ids are ignored. + # Bound ids to [0, vocab) before indexing logits: MLX does no bounds + # checking and out-of-bounds indexing is undefined behavior (crash / + # corruption), unlike torch's harmless negative wrap. MLX also lacks + # boolean-mask filtering, so out-of-range/negative ids route to a + # scratch slot at index vocab (dropped before the subtract) that never + # collides with a real token: real ids (including 0) are penalized + # once, strays ignored. valid = (generated >= 0) & (generated < vocab) safe = mx.where(valid, generated, vocab).astype(mx.int32) - # Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate - # ids are idempotent, so presence applies once per distinct token; the - # scratch column is discarded and the full-width subtract stays on-device. + # Scatter penalty into a (vocab + 1)-wide mask: duplicate ids are + # idempotent (presence applies once per token); scratch column dropped. mask = mx.zeros((vocab + 1,), dtype = logits.dtype) mask[safe] = penalty logits = logits - mask[:vocab] @@ -93,7 +131,7 @@ class MLXInferenceBackend: self.loaded_local_models = [] self.device = "mlx" self._generation_lock = threading.Lock() - # usage/timings of the latest generation; shipped on gen_done. + # usage/timings of the latest generation, shipped on gen_done. self.last_generation_stats = None self._model = None @@ -101,6 +139,9 @@ class MLXInferenceBackend: self._processor = None self._is_vlm = False self._config = {} + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 # Recorded for unload to release pinned memory back to the OS. self._memory_limits_applied = {} @@ -145,19 +186,26 @@ class MLXInferenceBackend: trust_remote_code = False, gpu_ids = None, dtype = None, + parallel_mode = None, + distributed_group = None, ) -> bool: import mlx.core as mx - # Keep the token so the native-template fallback can fetch a - # gated model's repo template later during generation. + # Keep the token so the native-template fallback can fetch a gated + # model's repo template during generation. self._hf_token = hf_token model_name = config.identifier if hasattr(config, "identifier") else str(config) is_vision = getattr(config, "is_vision", False) + distributed_rank, distributed_size = _mlx_distributed_rank_size(distributed_group) + is_distributed = distributed_group is not None and distributed_size > 1 + self._distributed_group = distributed_group + self._distributed_rank = distributed_rank + self._distributed_world_size = distributed_size - # GGUF guard. GGUF models are served by llama-server in the parent - # process, not mlx-lm here. Reaching this with is_gguf=True means the - # route's first detection flaked (transient HF Hub) but the subprocess - # re-detected GGUF; raise loudly instead of a cryptic mlx_lm error. + # GGUF guard: GGUF is served by llama-server in the parent process, + # not mlx-lm. Reaching here with is_gguf=True means the route's + # detection flaked but the subprocess re-detected GGUF; raise loudly + # instead of a cryptic mlx_lm error. if getattr(config, "is_gguf", False): raise RuntimeError( f"MLXInferenceBackend cannot load GGUF model '{model_name}': " @@ -176,11 +224,26 @@ class MLXInferenceBackend: is_lora = getattr(config, "is_lora", False) logger.info( - "Loading %s via %s (is_lora=%s)", + "Loading %s via %s (is_lora=%s, distributed=%s, rank=%s/%s, mode=%s)", model_name, "mlx-vlm" if is_vision else "mlx-lm", is_lora, + is_distributed, + distributed_rank, + distributed_size, + parallel_mode, ) + if is_distributed and parallel_mode not in ("pipeline", "tensor"): + raise ValueError( + "Unsloth: distributed MLX inference requires parallel_mode='pipeline' " + "or parallel_mode='tensor'." + ) + if is_distributed and is_lora: + raise ValueError( + "Unsloth: distributed MLX inference for LoRA adapter repos " + "is not supported yet. Merge/export the adapter into an MLX model " + "before distributed inference." + ) try: from unsloth_zoo.mlx.loader import FastMLXModel @@ -190,14 +253,23 @@ class MLXInferenceBackend: "(unsloth_zoo.mlx.loader). Reinstall via install.sh on Apple Silicon." ) from e + load_kwargs = { + "max_seq_length": max_seq_length, + "dtype": dtype, + "load_in_4bit": load_in_4bit, + "token": hf_token, + "trust_remote_code": trust_remote_code, + "text_only": False if is_vision else True, + } + if is_distributed: + if parallel_mode == "pipeline": + load_kwargs["pipeline_group"] = distributed_group + else: + load_kwargs["tensor_group"] = distributed_group + model, tokenizer_or_processor = FastMLXModel.from_pretrained( model_name, - max_seq_length = max_seq_length, - dtype = dtype, - load_in_4bit = load_in_4bit, - token = hf_token, - trust_remote_code = trust_remote_code, - text_only = False if is_vision else True, + **load_kwargs, ) if is_vision: @@ -217,8 +289,7 @@ class MLXInferenceBackend: self.models[model_name] = { # Per-model token for the native-template fallback (matches transformers). "hf_token": hf_token, - # Per-model consent for the native-template reload: re-use the exact - # trust_remote_code this model was loaded with (matches transformers). + # Per-model trust_remote_code reused by the native-template reload (matches transformers). "trust_remote_code": trust_remote_code, "model": self._model, "tokenizer": self._tokenizer, @@ -234,8 +305,7 @@ class MLXInferenceBackend: "has_audio_input": False, "context_length": runtime_context_length(self._model, max_seq_length), } - # Capture chat_template_info so the worker IPC reply ships it back and - # the route layer classifies capabilities like the other paths. + # Capture chat_template_info for the worker IPC reply and route capability classification. self._populate_chat_template_info(model_name) logger.info("Model %s loaded successfully", model_name) @@ -293,6 +363,9 @@ class MLXInferenceBackend: self._model = None self._tokenizer = None self._processor = None + self._distributed_group = None + self._distributed_rank = 0 + self._distributed_world_size = 1 if self.active_model_name == model_name: self.active_model_name = None gc.collect() @@ -320,8 +393,7 @@ class MLXInferenceBackend: max_new_tokens = 256, repetition_penalty = 1.0, cancel_event = None, - # Reasoning / tool kwargs forwarded by the route + worker; rendered via - # apply_chat_template_for_generation like the transformers path. + # Reasoning / tool kwargs, rendered via apply_chat_template_for_generation (transformers parity). tools = None, enable_thinking = None, reasoning_effort = None, @@ -334,7 +406,6 @@ class MLXInferenceBackend: # Reset so a failed run cannot surface stale stats. self.last_generation_stats = None - # Build messages with system prompt full_messages = [] if system_prompt: full_messages.append({"role": "system", "content": system_prompt}) @@ -351,7 +422,6 @@ class MLXInferenceBackend: {"type": "text", "text": content}, ] elif isinstance(content, list): - # Prepend image if not already present has_image = any( p.get("type") == "image" for p in content if isinstance(p, dict) ) @@ -429,11 +499,11 @@ class MLXInferenceBackend: if prompt is None: raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible") - # Same parity fix as the transformers backend: if the template dropped the - # requested tools, fall back to the native template so MLX text models keep - # advertising them. ``self._tokenizer`` is this entry's model_info tokenizer, - # so probe and native render share a renderer. (The VLM path renders via the - # processor for image tokens and is intentionally not wired here.) + # Parity with the transformers backend: if the template dropped the + # requested tools, fall back to the native template so MLX text models + # keep advertising them. self._tokenizer is this entry's tokenizer, so + # probe and native render share a renderer. (VLM renders via the + # processor for image tokens and is not wired here.) model_info = self.models.get(self.active_model_name, {}) prompt = render_with_native_template_fallback( formatted_prompt = prompt, @@ -455,7 +525,7 @@ class MLXInferenceBackend: min_p = float(min_p or 0.0), min_tokens_to_keep = 1, ) - # Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths). + # Repetition and/or presence penalty processors (GGUF/safetensors parity). logits_processors = [] if repetition_penalty is not None and float(repetition_penalty) not in ( 0.0, @@ -496,7 +566,6 @@ class MLXInferenceBackend: ): final_response = response token_ids.append(response.token) - # Decode full sequence with skip_special_tokens cumulative = self._tokenizer.decode( token_ids, skip_special_tokens = True, @@ -544,8 +613,7 @@ class MLXInferenceBackend: ) # Pick the chat-template-aware caller: processors with their own - # apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it - # directly; else fall back to the nested tokenizer. + # apply_chat_template + chat_template (e.g. Qwen2.5-VL), else the nested tokenizer. chat_target = self._processor if ( getattr(self._processor, "apply_chat_template", None) is None @@ -572,10 +640,9 @@ class MLXInferenceBackend: len(prompt), image is not None, ) - # mlx_vlm.stream_generate forwards **kwargs into generate_step, which - # builds the sampler + logits_processors internally. - # GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=`` - # silently falls into **kwargs and is ignored, stuck at greedy 0.0. + # stream_generate forwards **kwargs into generate_step (builds the + # sampler + logits_processors internally). GOTCHA: generate_step expects + # temperature= (long form); temp= is silently ignored, stuck at greedy 0.0. vlm_kwargs = dict( max_tokens = max_new_tokens, temperature = temperature, @@ -589,7 +656,7 @@ class MLXInferenceBackend: ) if presence_penalty: # Presence needs a custom processor: pass the full list (repetition + - # presence) instead of the repetition_penalty shortcut so both apply once. + # presence) instead of the repetition_penalty shortcut so both apply. from mlx_lm.sample_utils import make_logits_processors _vlm_processors = [] @@ -634,7 +701,7 @@ class MLXInferenceBackend: cancel_event = None, **gen_kwargs, ) -> Generator[str, None, None]: - # MLX LoRA adapter toggling not yet supported — generate normally + # MLX LoRA adapter toggling not yet supported; generate normally yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs) def reset_generation_state(self): diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py index cf5d24c367..4fa0d3ed26 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -50,6 +50,18 @@ _DISPATCH_DRAIN_TIMEOUT = 5.0 _UNLOAD_GEN_LOCK_TIMEOUT = 15.0 +class GenStreamError(str): + """A stream chunk carrying a real backend/generation error, not model text. + + Subclasses str so existing display/logging consumers are unaffected, while + callers that must abort a distributed run on error (raise_on_streamed_error) + can distinguish a real error from model output whose visible text starts with + "Error:" by checking isinstance(chunk, GenStreamError). + """ + + __slots__ = () + + class InferenceOrchestrator: """ Inference backend orchestrator — subprocess-based. @@ -482,13 +494,13 @@ class InferenceOrchestrator: initial_resp_queue = self._resp_queue while True: if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue: - yield f"Error: {self._subprocess_crash_message(crash_context)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return resp = read_one(read_timeout) if resp is None: # Check subprocess health if not self._ensure_subprocess_alive(): - yield f"Error: {self._subprocess_crash_message(crash_context)}" + yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}") return continue @@ -498,7 +510,7 @@ class InferenceOrchestrator: # Subprocess-level error (no request_id); request-scoped failures # arrive as gen_error below. if rtype == "error" and not resp.get("request_id"): - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return if rtype == "token": @@ -513,7 +525,7 @@ class InferenceOrchestrator: stats_holder["stats"] = resp.get("stats") return elif rtype == "gen_error": - yield f"Error: {resp.get('error', 'Unknown error')}" + yield GenStreamError(f"Error: {resp.get('error', 'Unknown error')}") return # ------------------------------------------------------------------ @@ -640,11 +652,11 @@ class InferenceOrchestrator: GPU work stays serialized; this only avoids orchestrator lock contention. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return # Latch the target model so the recheck below can detect a switch that completed # between _start_dispatcher and mailbox registration (mirrors the locked path's @@ -655,7 +667,7 @@ class InferenceOrchestrator: # so without this early-out a compare request would enqueue a generate on the # outgoing model and delay the switch. if self._unload_pending: - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return # Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under @@ -727,7 +739,7 @@ class InferenceOrchestrator: # _stop_dispatcher joins the dispatcher, which itself takes that lock. if orphaned_dispatcher: self._stop_dispatcher() - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return try: @@ -735,7 +747,7 @@ class InferenceOrchestrator: except RuntimeError as exc: with self._mailbox_lock: self._mailboxes.pop(request_id, None) - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return def read_mailbox(timeout): @@ -813,6 +825,59 @@ class InferenceOrchestrator: self._stop_dispatcher() return True + def share_distributed_object( + self, + obj, + timeout: Optional[float] = 300.0, + ): + """Share a small object through the worker's MLX distributed group.""" + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess is not running") + + self._wait_dispatcher_idle() + with self._mailbox_lock: + if self._mailboxes: + raise RuntimeError( + "Cannot share distributed objects while compare requests are active" + ) + request_id = str(uuid.uuid4()) + cmd = { + "type": "share_object", + "request_id": request_id, + "object": obj, + } + + with self._gen_lock: + self._send_cmd(cmd) + deadline = None if timeout is None else time.monotonic() + timeout + while deadline is None or time.monotonic() < deadline: + remaining = 1.0 if deadline is None else max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout = min(remaining, 1.0)) + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError(self._subprocess_crash_message("sharing chat turn")) + continue + + rtype = resp.get("type", "") + rid = resp.get("request_id") + if rid and rid != request_id: + logger.debug( + "Skipping response for request_id=%s while sharing request_id=%s", + rid, + request_id, + ) + continue + if rtype == "shared": + return resp.get("object") + if rtype == "share_error": + raise RuntimeError(resp.get("error", "Failed to share object")) + if rtype == "error": + raise RuntimeError(resp.get("error", "Subprocess error")) + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for distributed object share") + # ------------------------------------------------------------------ # Public API — same interface as InferenceBackend # ------------------------------------------------------------------ @@ -828,6 +893,8 @@ class InferenceOrchestrator: approved_remote_code_fingerprint: Optional[str] = None, gpu_ids: Optional[list[int]] = None, subject: Optional[str] = None, + tensor_parallel: bool = False, + mlx_distributed: bool = False, ) -> bool: """Load a model for inference. @@ -853,6 +920,11 @@ class InferenceOrchestrator: "approved_remote_code_fingerprint": approved_remote_code_fingerprint, "subject": subject, "gpu_ids": gpu_ids, + "tensor_parallel": bool(tensor_parallel), + "mlx_distributed": bool(mlx_distributed), + "mlx_parallel_mode": ("tensor" if tensor_parallel else "pipeline") + if mlx_distributed + else None, } resolved_gpu_ids, gpu_selection = prepare_gpu_selection( gpu_ids, @@ -1338,11 +1410,11 @@ class InferenceOrchestrator: readers don't consume each other's tokens off the shared resp_queue. """ if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return expected_model = self.active_model_name @@ -1359,7 +1431,7 @@ class InferenceOrchestrator: # so we never generate on the wrong one. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return request_id = str(uuid.uuid4()) image_b64 = self._pil_to_base64(image) if image is not None else None @@ -1385,7 +1457,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( @@ -1544,10 +1616,10 @@ class InferenceOrchestrator: ) -> Generator[str, None, None]: """Shared inner logic for audio input generation (Whisper + ASR).""" if not self._ensure_subprocess_alive(): - yield "Error: Inference subprocess is not running" + yield GenStreamError("Error: Inference subprocess is not running") return if not self.active_model_name: - yield "Error: No active model" + yield GenStreamError("Error: No active model") return expected_model = self.active_model_name @@ -1556,7 +1628,7 @@ class InferenceOrchestrator: # cleared or swapped the model while we waited. if self._unload_pending or self.active_model_name != expected_model: # Won the lock handoff during a switch; don't start on the outgoing model. - yield "Error: model is being unloaded" + yield GenStreamError("Error: model is being unloaded") return request_id = str(uuid.uuid4()) @@ -1583,7 +1655,7 @@ class InferenceOrchestrator: try: self._send_cmd(cmd) except RuntimeError as exc: - yield f"Error: {exc}" + yield GenStreamError(f"Error: {exc}") return yield from self._consume_token_stream( diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index d4b102e422..05dee39283 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -13,6 +13,7 @@ mp.Queue, and exits on shutdown or unload. Pattern follows core/training/worker. from __future__ import annotations import base64 +import json from loggers import get_logger import os import queue as _queue @@ -26,6 +27,9 @@ from typing import Any logger = get_logger(__name__) from utils.hardware import apply_gpu_ids +_SHARE_OBJECT_MAX_BYTES = 1 << 20 +_SHARE_OBJECT_ERROR_SIZE = -1 + # studio/backend root, prepended to sys.path so the spawned subprocess can # import the utils/core packages. _BACKEND_PATH = str(Path(__file__).resolve().parent.parent.parent) @@ -75,6 +79,17 @@ def _send_response(resp_queue: Any, response: dict) -> None: logger.error("Failed to send response: %s", exc) +def _encode_share_object(obj: Any) -> bytes: + data = json.dumps(obj, separators = (",", ":"), ensure_ascii = False).encode("utf-8") + if len(data) > _SHARE_OBJECT_MAX_BYTES: + raise ValueError("Distributed object share payload is too large") + return data + + +def _decode_share_object(data: Any) -> Any: + return json.loads(bytes(data.tolist()).decode("utf-8")) + + def _clean_token(value: str | None) -> str | None: """Normalize an HF token: blank or whitespace-only becomes None.""" return value if value and value.strip() else None @@ -329,14 +344,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1", ) try: - success = backend.load_model( - config = mc, - max_seq_length = config.get("max_seq_length", 2048), - load_in_4bit = load_in_4bit, - hf_token = hf_token, - trust_remote_code = trust_remote_code, - gpu_ids = config.get("resolved_gpu_ids"), - ) + load_kwargs = { + "config": mc, + "max_seq_length": config.get("max_seq_length", 2048), + "load_in_4bit": load_in_4bit, + "hf_token": hf_token, + "trust_remote_code": trust_remote_code, + "gpu_ids": config.get("resolved_gpu_ids"), + } + if getattr(backend, "device", None) == "mlx": + load_kwargs["parallel_mode"] = config.get("mlx_parallel_mode") + load_kwargs["distributed_group"] = config.get("_mlx_distributed_group") + success = backend.load_model(**load_kwargs) finally: heartbeat_stop.set() @@ -521,6 +540,67 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None: ) +def _handle_share_object(backend, cmd: dict, resp_queue: Any) -> None: + """Share a small Python object across MLX distributed ranks.""" + request_id = cmd.get("request_id", "") + group = getattr(backend, "_distributed_group", None) + rank = int(getattr(backend, "_distributed_rank", 0) or 0) + world_size = int(getattr(backend, "_distributed_world_size", 1) or 1) + obj = cmd.get("object") + + try: + if group is None or world_size <= 1: + shared = obj + else: + import mlx.core as mx + if rank == 0: + if obj is None: + mx.eval(mx.distributed.all_sum(mx.array(0), group = group)) + shared = None + else: + try: + data = mx.array(_encode_share_object(obj), dtype = mx.uint8) + except Exception: + mx.eval( + mx.distributed.all_sum( + mx.array(_SHARE_OBJECT_ERROR_SIZE), + group = group, + ) + ) + raise + mx.eval(mx.distributed.all_sum(mx.array(data.size), group = group)) + mx.eval(mx.distributed.all_sum(data, group = group)) + shared = obj + else: + size = int(mx.distributed.all_sum(mx.array(0), group = group).item()) + if size == _SHARE_OBJECT_ERROR_SIZE: + raise RuntimeError("Failed to share distributed object") + if size == 0: + shared = None + else: + data = mx.zeros(size, dtype = mx.uint8) + data = mx.distributed.all_sum(data, group = group) + shared = _decode_share_object(data) + _send_response( + resp_queue, + { + "type": "shared", + "request_id": request_id, + "object": shared, + }, + ) + except Exception as exc: + _send_response( + resp_queue, + { + "type": "share_error", + "request_id": request_id, + "error": str(exc), + "stack": traceback.format_exc(limit = 20), + }, + ) + + def _handle_generate_audio(backend, cmd: dict, resp_queue: Any) -> None: """Handle TTS audio generation — returns WAV bytes + sample_rate.""" request_id = cmd.get("request_id", "") @@ -720,9 +800,29 @@ def run_inference_process( exc, ) try: - from core.inference.mlx_inference import MLXInferenceBackend + from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed backend = MLXInferenceBackend() + if config.get("mlx_distributed"): + group, rank, size = _init_mlx_distributed() + config["_mlx_distributed_group"] = group + if size <= 1: + # A singleton group (MLX built without distributed support, + # or an invalid launch env/hostfile) would leave nonzero ranks + # looping forever on share_distributed_object. Fail the load + # instead of silently continuing without sharding. + raise RuntimeError( + "MLX distributed launch requested but initialized a singleton " + "group (size 1). Ensure the installed MLX has distributed " + "support and the launch environment/hostfile is valid, or run " + "without distributed." + ) + logger.info( + "MLX distributed initialized in worker: rank=%s size=%s mode=%s", + rank, + size, + config.get("mlx_parallel_mode"), + ) _send_response( resp_queue, {"type": "status", "message": "Loading model..."}, @@ -764,6 +864,8 @@ def run_inference_process( if _drain_skip_generate(cmd, resp_queue, drain_event): continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) @@ -977,6 +1079,9 @@ def run_inference_process( continue _handle_generate(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "share_object": + _handle_share_object(backend, cmd, resp_queue) + elif cmd_type == "load": if backend.active_model_name: backend.unload_model(backend.active_model_name) diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index ce755ae1fb..d9901a5b2e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -40,6 +40,43 @@ def _positive_int_or_none(value: Any) -> Optional[int]: return value_int if value_int > 0 else None +def _nonnegative_int_or_none(value: Any) -> Optional[int]: + if isinstance(value, bool): + return None + try: + value_int = int(value) + except (TypeError, ValueError): + return None + return value_int if value_int >= 0 else None + + +_MLX_MPI_DISTRIBUTED_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) + + +def _mlx_distributed_launch_detected() -> bool: + if _nonnegative_int_or_none(os.environ.get("MLX_RANK")) is not None: + world_size = _positive_int_or_none(os.environ.get("MLX_WORLD_SIZE")) + if world_size is not None and world_size > 1: + return True + return bool( + os.environ.get("MLX_HOSTFILE") + or os.environ.get("MLX_IBV_DEVICES") + or os.environ.get("MLX_JACCL_COORDINATOR") + or (os.environ.get("NCCL_HOST_IP") and os.environ.get("NCCL_PORT")) + ) + return any( + _nonnegative_int_or_none(os.environ.get(rank_env)) is not None + and (_positive_int_or_none(os.environ.get(size_env)) or 0) > 1 + for rank_env, size_env in _MLX_MPI_DISTRIBUTED_ENV_PAIRS + ) + + def _install_httpcore_asyncgen_silencer() -> None: """Silence benign httpx/httpcore asyncgen GC noise on Python 3.13. @@ -3426,6 +3463,15 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre status_code = 400, detail = "gpu_ids is not supported for GGUF models yet.", ) + if not config.is_gguf and _mlx_distributed_launch_detected(): + raise HTTPException( + status_code = 400, + detail = ( + "Studio does not support distributed MLX inference under " + "mlx.launch. Use `mlx.launch ... unsloth chat` or run Studio " + "without the distributed launcher." + ), + ) # Effective quantization (LoRA can flip 4-bit -> 16-bit); guard + load reuse it. effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit) diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index ac4088fb25..55a3198a6b 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -4,6 +4,8 @@ import sys import types from types import SimpleNamespace +import pytest + class _DummyMetal: @staticmethod @@ -185,6 +187,129 @@ def test_mlx_inference_vlm_lora_uses_unsloth_loader_without_native_adapter_rewri assert isinstance(backend._tokenizer, _DummyTokenizer) +def test_mlx_inference_distributed_vlm_forwards_group_to_fast_mlx(monkeypatch): + _install_fake_mlx(monkeypatch) + calls = [] + _install_fake_fast_mlx(monkeypatch, calls) + from core.inference.mlx_inference import MLXInferenceBackend + + group = SimpleNamespace(size = lambda: 2, rank = lambda: 0) + config = SimpleNamespace(identifier = "fake/vlm", is_vision = True, is_lora = False) + for mode, group_key in (("tensor", "tensor_group"), ("pipeline", "pipeline_group")): + calls.clear() + assert MLXInferenceBackend().load_model(config, parallel_mode = mode, distributed_group = group) + _, kwargs = calls.pop() + assert kwargs["text_only"] is False and kwargs[group_key] is group + + calls.clear() + singleton = SimpleNamespace(size = lambda: 1, rank = lambda: 0) + assert MLXInferenceBackend().load_model( + config, parallel_mode = "tensor", distributed_group = singleton + ) + assert not {"tensor_group", "pipeline_group"} & set(calls.pop()[1]) + + config = SimpleNamespace(identifier = "fake/adapter", is_vision = False, is_lora = True) + with pytest.raises(ValueError, match = "LoRA adapter repos"): + MLXInferenceBackend().load_model(config, parallel_mode = "tensor", distributed_group = group) + + +@pytest.mark.parametrize("accepts_backend", (True, False)) +def test_mlx_distributed_init_selects_jaccl_backend(monkeypatch, accepts_backend): + _install_fake_mlx(monkeypatch) + from core.inference.mlx_inference import _init_mlx_distributed + + group = SimpleNamespace(rank = lambda: 1, size = lambda: 2) + calls = [] + + def _init(**kwargs): + calls.append(kwargs) + if kwargs and not accepts_backend: + raise TypeError("backend keyword unsupported") + return group + + sys.modules["mlx.core"].distributed = SimpleNamespace(init = _init) + monkeypatch.setenv("MLX_JACCL_COORDINATOR", "127.0.0.1:12345") + monkeypatch.setenv("MLX_IBV_DEVICES", "/tmp/devices.json") + + assert _init_mlx_distributed() == (group, 1, 2) + assert calls == ([{"backend": "jaccl"}] if accepts_backend else [{"backend": "jaccl"}, {}]) + + +def test_worker_share_object_receives_distributed_payload(monkeypatch): + from core.inference import worker + + shared_obj = {"type": "turn", "text": "hi"} + payload = worker._encode_share_object(shared_obj) + + def _array(value): + val = value.item() if hasattr(value, "item") else value + return SimpleNamespace( + item = lambda: val, + tolist = lambda: list(val) if hasattr(val, "__iter__") else [val], + ) + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.uint8 = "uint8" + mlx_core.array = _array + mlx_core.zeros = lambda *_a, **_k: _array([]) + + def _all_sum(value, group = None): + value = value.item() if hasattr(value, "item") else value + return _array(len(payload)) if value == 0 else _array(payload) + + mlx_core.distributed = SimpleNamespace(all_sum = _all_sum) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 1, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": None}, + SimpleNamespace(put = responses.append), + ) + + response = responses[0] + assert response["object"] == shared_obj + + +def test_worker_share_object_oversize_notifies_peers(monkeypatch): + from core.inference import worker + + calls = [] + + mlx_pkg = types.ModuleType("mlx") + mlx_core = types.ModuleType("mlx.core") + mlx_core.array = lambda value, **_kwargs: SimpleNamespace(item = lambda: value) + mlx_core.eval = lambda value: value + mlx_core.distributed = SimpleNamespace( + all_sum = lambda value, group = None: calls.append(value.item()) or value + ) + mlx_pkg.core = mlx_core + monkeypatch.setitem(sys.modules, "mlx", mlx_pkg) + monkeypatch.setitem(sys.modules, "mlx.core", mlx_core) + monkeypatch.setattr(worker, "_SHARE_OBJECT_MAX_BYTES", 8) + + responses = [] + worker._handle_share_object( + SimpleNamespace( + _distributed_group = object(), + _distributed_rank = 0, + _distributed_world_size = 2, + ), + {"type": "share_object", "request_id": "rid", "object": {"text": "too long"}}, + SimpleNamespace(put = responses.append), + ) + + assert calls == [worker._SHARE_OBJECT_ERROR_SIZE] + assert responses[0]["type"] == "share_error" + + # Regression: generate_chat_response must accept the four template kwargs # (tools / enable_thinking / reasoning_effort / preserve_thinking) so the route # layer can forward UI toggles. The old signature raised diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index e8a3414d73..c3b188710e 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -4,9 +4,11 @@ """Model loading and streaming shared by `inference` and `chat`.""" import asyncio +import json import os import re import sys +from contextlib import contextmanager, redirect_stderr, redirect_stdout from pathlib import Path from typing import List, Optional @@ -14,10 +16,18 @@ import typer _THINK_OPEN = "" _THINK_BLOCK = re.compile(rf"{re.escape(_THINK_OPEN)}.*?", re.DOTALL) +_STREAMED_ERROR_PREFIX = "Error: " # Cloudflare (in front of remote Studio proxies like RunPod) 403s the default # "Python-urllib/X.Y" User-Agent as a bot; send a real one on every request. _USER_AGENT = "unsloth-cli" +_MPI_ENV_PAIRS = ( + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +) # Built lazily; urllib stays function-local to match this module. _no_redirect_opener = None @@ -61,6 +71,108 @@ def configure_quiet_logging() -> None: os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") +def _parse_nonnegative_int(value: Optional[str]) -> Optional[int]: + if value is None: + return None + try: + parsed = int(value) + except (TypeError, ValueError): + return None + return parsed if parsed >= 0 else None + + +def _first_mpi_env_pair() -> tuple[Optional[int], Optional[int]]: + for rank_name, size_name in _MPI_ENV_PAIRS: + rank = _parse_nonnegative_int(os.environ.get(rank_name)) + world_size = _parse_nonnegative_int(os.environ.get(size_name)) + if rank is not None and world_size is not None and world_size > 1 and rank < world_size: + return rank, world_size + return None, None + + +def _json_rank_count_from_env(name: str) -> Optional[int]: + value = os.environ.get(name) + if not value: + return None + try: + if value.lstrip().startswith(("[", "{")): + data = json.loads(value) + else: + with open(value, "r") as f: + data = json.load(f) + except (OSError, json.JSONDecodeError): + return None + if isinstance(data, list): + return len(data) + if isinstance(data, dict) and isinstance(data.get("hosts"), list): + return len(data["hosts"]) + return None + + +def mlx_distributed_info() -> tuple[bool, int, Optional[int]]: + """Return launch-context metadata without initializing MLX distributed.""" + rank = _parse_nonnegative_int(os.environ.get("MLX_RANK")) + world_size = _parse_nonnegative_int(os.environ.get("MLX_WORLD_SIZE")) + if rank is not None: + if ( + world_size is not None + and world_size > 1 + and rank < world_size + and os.environ.get("NCCL_HOST_IP") + and os.environ.get("NCCL_PORT") + ): + return True, rank, world_size + inferred_size = _json_rank_count_from_env("MLX_HOSTFILE") + if inferred_size is not None and inferred_size > 1 and rank < inferred_size: + return True, rank, inferred_size + inferred_size = _json_rank_count_from_env("MLX_IBV_DEVICES") + if ( + inferred_size is not None + and inferred_size > 1 + and rank < inferred_size + and os.environ.get("MLX_JACCL_COORDINATOR") + ): + return True, rank, inferred_size + return False, 0, None + + mpi_rank, mpi_world_size = _first_mpi_env_pair() + return mpi_rank is not None, mpi_rank or 0, mpi_world_size + + +def mlx_distributed_uses_mpi() -> bool: + """Whether the current distributed context was launched through MPI.""" + return ( + _parse_nonnegative_int(os.environ.get("MLX_RANK")) is None + and _first_mpi_env_pair()[0] is not None + ) + + +@contextmanager +def quiet_if_nonzero_mlx_rank(): + """Silence parent and child-process stdout/stderr on nonzero ranks.""" + if mlx_distributed_info()[1] == 0: + yield + return + + sys.stdout.flush() + sys.stderr.flush() + saved_stdout_fd = os.dup(1) + saved_stderr_fd = os.dup(2) + with open(os.devnull, "w") as devnull: + try: + os.dup2(devnull.fileno(), 1) + os.dup2(devnull.fileno(), 2) + with redirect_stdout(devnull), redirect_stderr(devnull): + yield + finally: + sys.stdout.flush() + sys.stderr.flush() + os.dup2(saved_stdout_fd, 1) + os.dup2(saved_stderr_fd, 2) + os.close(saved_stdout_fd) + os.close(saved_stderr_fd) + + def visible_text(text: str, show_thinking: bool) -> str: if show_thinking: return text @@ -120,6 +232,21 @@ def collect_stream(stream, show_thinking: bool) -> str: return visible_text(raw, show_thinking) +def raise_on_streamed_error(stream): + # Match real backend errors by type (GenStreamError), not the "Error:" text + # prefix, so a completion whose text opens with "Error:" is not misread as a + # failure that aborts a distributed run. + try: + ensure_studio_backend_path() + from core.inference.orchestrator import GenStreamError + except Exception: + GenStreamError = None + for chunk in stream: + if GenStreamError is not None and isinstance(chunk, GenStreamError): + raise RuntimeError(str(chunk)[len(_STREAMED_ERROR_PREFIX) :].strip() or "Unknown error") + yield chunk + + def render_columns( left_label: str, left_text: str, @@ -200,6 +327,19 @@ class ChatBackend: except Exception: pass + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + if self._kind != "unsloth" or not hasattr(self._backend, "share_distributed_object"): + raise RuntimeError( + "Distributed MLX chat requires the Unsloth MLX backend; " + f"backend '{self._kind}' cannot broadcast chat turns." + ) + return self._backend.share_distributed_object(obj, timeout = timeout) + def resolve_model_config(model: str, *, hf_token: Optional[str]): ensure_studio_backend_path() @@ -293,36 +433,59 @@ def load_chat_backend( fresh_backend uses a private orchestrator so a second model (compare's base column) can run alongside the main one. """ - if model_config is None: - model_config = resolve_model_config(model, hf_token = hf_token) + with quiet_if_nonzero_mlx_rank(): + is_mlx_distributed, rank, _world_size = mlx_distributed_info() + if model_config is None: + model_config = resolve_model_config(model, hf_token = hf_token) - typer.echo(f"Loading {model}", err = True) + if is_mlx_distributed and model_config.is_gguf: + if rank == 0: + typer.echo( + "Distributed MLX inference does not support GGUF/llama.cpp models. " + "Use a non-GGUF MLX model under mlx.launch, or run GGUF without " + "mlx.launch.", + err = True, + ) + raise typer.Exit(code = 1) - if model_config.is_gguf: - return _load_gguf_backend( - model_config, - hf_token = hf_token, - max_seq_length = max_seq_length, - tensor_parallel = tensor_parallel, - llama_extra_args = llama_extra_args, - ) + if rank == 0: + typer.echo(f"Loading {model}", err = True) - if fresh_backend: - ensure_studio_backend_path() - from core.inference import InferenceOrchestrator - backend = InferenceOrchestrator() - else: - ensure_studio_backend_path() - from core.inference import get_inference_backend - backend = get_inference_backend() - if not backend.load_model( - config = model_config, - max_seq_length = max_seq_length, - load_in_4bit = load_in_4bit, - hf_token = hf_token, - ): - typer.echo("Model load failed", err = True) - raise typer.Exit(code = 1) + if model_config.is_gguf: + return _load_gguf_backend( + model_config, + hf_token = hf_token, + max_seq_length = max_seq_length, + tensor_parallel = tensor_parallel, + llama_extra_args = llama_extra_args, + ) + + if fresh_backend: + ensure_studio_backend_path() + from core.inference import InferenceOrchestrator + backend = InferenceOrchestrator() + else: + ensure_studio_backend_path() + from core.inference import get_inference_backend + backend = get_inference_backend() + try: + loaded = backend.load_model( + config = model_config, + max_seq_length = max_seq_length, + load_in_4bit = load_in_4bit, + hf_token = hf_token, + tensor_parallel = tensor_parallel, + mlx_distributed = is_mlx_distributed, + ) + except Exception as exc: + if not is_mlx_distributed: + raise + if rank == 0: + typer.echo(str(exc) or "Model load failed", err = True) + raise typer.Exit(code = 1) + if not loaded: + typer.echo("Model load failed", err = True) + raise typer.Exit(code = 1) return ChatBackend("unsloth", backend) diff --git a/unsloth_cli/commands/chat.py b/unsloth_cli/commands/chat.py index d3bfbbf96b..bc4a72f36c 100644 --- a/unsloth_cli/commands/chat.py +++ b/unsloth_cli/commands/chat.py @@ -1,6 +1,7 @@ # 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 sys from typing import List, Optional import typer @@ -12,6 +13,10 @@ from unsloth_cli._inference import ( connect_studio_server, ensure_studio_backend_path, load_chat_backend, + mlx_distributed_info, + mlx_distributed_uses_mpi, + quiet_if_nonzero_mlx_rank, + raise_on_streamed_error, render_columns, resolve_model_config, stream_markdown, @@ -107,6 +112,20 @@ def _compare_needs_second_model() -> bool: return False +def _drain_available_stdin() -> None: + """Drain already-buffered launcher stdin on nonzero distributed ranks.""" + try: + import os + from select import select + + fd = sys.stdin.fileno() + while select([fd], [], [], 0)[0]: + if not os.read(fd, 8192): + break + except Exception: + return + + def _pick_trained_model(console) -> str: ensure_studio_backend_path() from utils.models import scan_trained_models @@ -158,7 +177,8 @@ def chat( "--tensor-parallel/--no-tensor-parallel", help = ( "Split a GGUF across GPUs by tensor (--split-mode tensor) instead " - "of by layer. Ignored for non-GGUF models." + "of by layer. Under non-MPI mlx.launch, select MLX tensor " + "parallel mode instead of pipeline mode." ), ), llama_extra_args: Optional[List[str]] = typer.Option( @@ -195,15 +215,43 @@ def chat( console = Console() err = Console(stderr = True) + is_mlx_distributed, rank, _world_size = mlx_distributed_info() + should_print = rank == 0 + + if is_mlx_distributed and mlx_distributed_uses_mpi(): + if should_print: + err.print( + "Distributed `unsloth chat` with MPI needs rank-0 prompt broadcast, " + "which is not enabled yet. Use a non-MPI MLX launcher backend " + "such as ring/JACCL for now.", + style = "red", + markup = False, + ) + raise typer.Exit(code = 1) if model is None: + if is_mlx_distributed: + if should_print: + err.print( + "Distributed `unsloth chat` requires an explicit model id or path.", + style = "red", + markup = False, + ) + raise typer.Exit(code = 1) model = _pick_trained_model(console) # Resolve first so --compare can be rejected before the slow load. - model_config = resolve_model_config(model, hf_token = hf_token) + with quiet_if_nonzero_mlx_rank(): + model_config = resolve_model_config(model, hf_token = hf_token) compare_blocked = _compare_blocked_reason(model_config) + if is_mlx_distributed: + compare_blocked = ( + "distributed MLX chat does not support compare mode yet because it " + "would need a second distributed worker group on the same ranks" + ) if compare and compare_blocked: - err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False) + if should_print: + err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False) raise typer.Exit(code = 1) load_opts = dict( @@ -215,9 +263,11 @@ def chat( ) # Prefer a running Studio server: instant starts, model shared with the UI. - chat_backend = None if no_server else connect_studio_server(model, **load_opts) + chat_backend = ( + None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts) + ) server_mode = chat_backend is not None - if server_mode: + if server_mode and should_print: console.print( "(Studio server connected — model stays warm after /exit)", style = "bright_black", @@ -242,23 +292,26 @@ def chat( return True base_id = model_config.base_model if not base_id: - console.print( - "(compare unavailable: this adapter doesn't record its base model)", - style = "yellow", - ) + if should_print: + console.print( + "(compare unavailable: this adapter doesn't record its base model)", + style = "yellow", + ) return False - console.print( - f"(loading base model {base_id} for compare — keeps two models in memory)", - style = "bright_black", - markup = False, - ) + if should_print: + console.print( + f"(loading base model {base_id} for compare — keeps two models in memory)", + style = "bright_black", + markup = False, + ) try: # Use the same precision as the tuned model for fair comparison base_load_opts = dict(load_opts) # Copy original options base_load_opts["load_in_4bit"] = _get_base_load_in_4bit(model_config) base_backend = load_chat_backend(base_id, fresh_backend = True, **base_load_opts) except Exception as exc: - err.print(f"(base model load failed: {exc})", style = "red", markup = False) + if should_print: + err.print(f"(base model load failed: {exc})", style = "red", markup = False) return False return True @@ -267,7 +320,7 @@ def chat( def generate(backend = None, use_adapter = None): # Reads messages and show_thinking live, so /reset and /think apply. - return (backend or chat_backend).stream( + stream = (backend or chat_backend).stream( messages, system_prompt = system_prompt, temperature = temperature, @@ -278,22 +331,48 @@ def chat( enable_thinking = show_thinking, use_adapter = use_adapter, ) + return raise_on_streamed_error(stream) if is_mlx_distributed else stream - console.print() - console.print(f"Chatting with {name}", style = "bold green", markup = False) - console.print(_HELP, style = "bright_black") + if should_print: + console.print() + console.print(f"Chatting with {name}", style = "bold green", markup = False) + console.print(_HELP, style = "bright_black") # legacy_windows: pre-VT consoles print raw ANSI as ←[1;36m garbage. - you_prompt = _you_prompt(console.is_terminal and not console.legacy_windows) + you_prompt = ( + _you_prompt(console.is_terminal and not console.legacy_windows) if should_print else "" + ) assistant_label = "[bold magenta]Assistant:[/bold magenta]" try: while True: - try: - user = input(you_prompt).strip() - except (EOFError, KeyboardInterrupt): - console.print() - break + if should_print: + try: + user = input(you_prompt).strip() + except (EOFError, KeyboardInterrupt): + if should_print: + console.print() + user = "/exit" + turn = {"type": "turn", "text": user} + else: + turn = None + + if is_mlx_distributed: + try: + turn = chat_backend.share_distributed_object(turn, timeout = None) + if not should_print: + _drain_available_stdin() + except Exception as exc: + if should_print: + err.print( + f"\n(error sharing chat turn: {exc})", + style = "red", + markup = False, + ) + raise typer.Exit(code = 1) + if not turn: + continue + user = str(turn.get("text", "")).strip() if not user: continue @@ -301,55 +380,69 @@ def chat( break if user == "/reset": messages = [] - console.print("(history cleared)", style = "bright_black") + if should_print: + console.print("(history cleared)", style = "bright_black") continue if user == "/think": show_thinking = not show_thinking - state = "on" if show_thinking else "off" - console.print(f"(thinking {state})", style = "bright_black") + if should_print: + state = "on" if show_thinking else "off" + console.print(f"(thinking {state})", style = "bright_black") continue if user == "/compare": if compare_blocked: - console.print(f"(compare unavailable: {compare_blocked})", style = "yellow") + if should_print: + console.print(f"(compare unavailable: {compare_blocked})", style = "yellow") continue if not compare_mode and dual_compare and not load_base_for_compare(): continue compare_mode = not compare_mode - state = "on" if compare_mode else "off" - console.print(f"(compare {state})", style = "bright_black") + if should_print: + state = "on" if compare_mode else "off" + console.print(f"(compare {state})", style = "bright_black") continue if user in ("/help", "/?"): - console.print(_HELP, style = "bright_black") + if should_print: + console.print(_HELP, style = "bright_black") continue messages.append({"role": "user", "content": user}) try: if compare_mode: - console.print("(comparing base vs tuned…)", style = "bright_black") + if should_print: + console.print("(comparing base vs tuned…)", style = "bright_black") if dual_compare: base_text = collect_stream(generate(backend = base_backend), show_thinking) tuned_text = collect_stream(generate(), show_thinking) else: base_text = collect_stream(generate(use_adapter = False), show_thinking) tuned_text = collect_stream(generate(use_adapter = True), show_thinking) - console.print() - render_columns( - "base", base_text, f"{name} (tuned)", tuned_text, console = console - ) + if should_print: + console.print() + render_columns( + "base", base_text, f"{name} (tuned)", tuned_text, console = console + ) # History continues as the tuned model; base is just the reference. answer = tuned_text else: - console.print(assistant_label) - answer = stream_markdown(generate(), show_thinking, console = console) + if should_print: + console.print(assistant_label) + answer = stream_markdown(generate(), show_thinking, console = console) + else: + answer = collect_stream(generate(), show_thinking) except KeyboardInterrupt: # Ctrl-C aborts this answer only; drop the unanswered turn. - console.print("\n(interrupted)", style = "bright_black") + if should_print: + console.print("\n(interrupted)", style = "bright_black") messages.pop() continue except Exception as exc: - err.print(f"\n(error: {exc})", style = "red", markup = False) + if should_print: + err.print(f"\n(error: {exc})", style = "red", markup = False) messages.pop() + if is_mlx_distributed: + raise typer.Exit(code = 1) continue messages.append( @@ -359,4 +452,5 @@ def chat( chat_backend.close() if base_backend is not None: base_backend.close() - err.print("\nBye.", style = "bright_black") + if should_print: + err.print("\nBye.", style = "bright_black") diff --git a/unsloth_cli/commands/inference.py b/unsloth_cli/commands/inference.py index 1401fa9e9e..84a126163e 100644 --- a/unsloth_cli/commands/inference.py +++ b/unsloth_cli/commands/inference.py @@ -6,9 +6,13 @@ from typing import List, Optional import typer from unsloth_cli._inference import ( + collect_stream, configure_quiet_logging, connect_studio_server, load_chat_backend, + mlx_distributed_info, + mlx_distributed_uses_mpi, + raise_on_streamed_error, stream_to_stdout, ) @@ -36,7 +40,8 @@ def inference( "--tensor-parallel/--no-tensor-parallel", help = ( "Split a GGUF across GPUs by tensor (--split-mode tensor) instead " - "of by layer. Ignored for non-GGUF models." + "of by layer. Under non-MPI mlx.launch, select MLX tensor " + "parallel mode instead of pipeline mode." ), ), llama_extra_args: Optional[List[str]] = typer.Option( @@ -69,8 +74,20 @@ def inference( if not verbose: configure_quiet_logging() - # A running Studio server keeps the model warm between runs, which is - # exactly what a one-shot command wants. + is_mlx_distributed, rank, _world_size = mlx_distributed_info() + if is_mlx_distributed and mlx_distributed_uses_mpi(): + if rank == 0: + typer.echo( + "Distributed `unsloth inference` with MPI is not supported by " + "the current subprocess backend. Use a non-MPI MLX launcher " + "backend such as ring/JACCL for now.", + err = True, + ) + raise typer.Exit(code = 1) + + # A running Studio server keeps the model warm between runs. Under + # mlx.launch, every rank must enter the local MLX path instead of rank 0 + # alone talking to a server. load_opts = dict( hf_token = hf_token, max_seq_length = max_seq_length, @@ -78,7 +95,9 @@ def inference( tensor_parallel = tensor_parallel, llama_extra_args = llama_extra_args, ) - chat_backend = None if no_server else connect_studio_server(model, **load_opts) + chat_backend = ( + None if (no_server or is_mlx_distributed) else connect_studio_server(model, **load_opts) + ) if chat_backend is None: chat_backend = load_chat_backend(model, **load_opts) try: @@ -92,7 +111,23 @@ def inference( repetition_penalty = repetition_penalty, enable_thinking = think, ) - typer.echo("Assistant:") - stream_to_stdout(stream, show_thinking = think) + if is_mlx_distributed: + stream = raise_on_streamed_error(stream) + if rank == 0: + typer.echo("Assistant:") + try: + stream_to_stdout(stream, show_thinking = think) + except RuntimeError as exc: + if not is_mlx_distributed: + raise + typer.echo(f"Error: {exc}", err = True) + raise typer.Exit(code = 1) + else: + try: + collect_stream(stream, show_thinking = think) + except RuntimeError: + if not is_mlx_distributed: + raise + raise typer.Exit(code = 1) finally: chat_backend.close() diff --git a/unsloth_cli/tests/test_inference_chat.py b/unsloth_cli/tests/test_inference_chat.py index 013b7c5a81..56633408fb 100644 --- a/unsloth_cli/tests/test_inference_chat.py +++ b/unsloth_cli/tests/test_inference_chat.py @@ -26,6 +26,8 @@ from unsloth_cli._inference import ( ChatBackend, HttpChatBackend, collect_stream, + mlx_distributed_info, + mlx_distributed_uses_mpi, render_columns, visible_text, ) @@ -39,6 +41,16 @@ class _FakeConfig: path = None +_EXPECTED_MPI_ENV_PAIRS = [ + ("OMPI_COMM_WORLD_RANK", "OMPI_COMM_WORLD_SIZE"), + ("PMI_RANK", "PMI_SIZE"), + ("PMIX_RANK", "PMIX_SIZE"), + ("MPI_RANK", "MPI_WORLD_SIZE"), + ("MV2_COMM_WORLD_RANK", "MV2_COMM_WORLD_SIZE"), +] +_IGNORED_DISTRIBUTED_ENV_PAIRS = [("SLURM_PROCID", "SLURM_NTASKS")] + + def _chat_app(): cli = typer.Typer() cli.command()(chatmod.chat) @@ -53,6 +65,39 @@ def _inference_app(): return cli +def _clear_mlx_distributed_env(monkeypatch): + for name in ( + "MLX_RANK", + "MLX_HOSTFILE", + "MLX_WORLD_SIZE", + "MLX_IBV_DEVICES", + "MLX_JACCL_COORDINATOR", + "NCCL_HOST_IP", + "NCCL_PORT", + *(rank for rank, _size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS), + *(size for _rank, size in _EXPECTED_MPI_ENV_PAIRS + _IGNORED_DISTRIBUTED_ENV_PAIRS), + ): + monkeypatch.delenv(name, raising = False) + + +def _set_mlx_nccl_env( + monkeypatch, + *, + rank: str = "0", + size: str = "2", +): + monkeypatch.setenv("MLX_RANK", rank) + monkeypatch.setenv("MLX_WORLD_SIZE", size) + monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1") + monkeypatch.setenv("NCCL_PORT", "12345") + + +@pytest.fixture(autouse = True) +def _isolate_mlx_distributed_env(monkeypatch): + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.delenv("HF_TOKEN", raising = False) + + def test_visible_text_passthrough_when_shown(): text = "reasoninganswer" assert visible_text(text, show_thinking = True) == text @@ -101,6 +146,46 @@ def test_inference_exposes_gguf_runtime_options(): assert "--llama-extra-arg" in (getattr(extra, "param_decls", None) or []) +def test_mlx_distributed_info_reads_launch_env(monkeypatch, tmp_path): + _clear_mlx_distributed_env(monkeypatch) + assert mlx_distributed_info() == (False, 0, None) + assert mlx_distributed_uses_mpi() is False + + monkeypatch.setenv("MLX_RANK", "1") + monkeypatch.setenv("MLX_WORLD_SIZE", "2") + assert mlx_distributed_info() == (False, 0, None) + monkeypatch.setenv("NCCL_HOST_IP", "127.0.0.1") + monkeypatch.setenv("NCCL_PORT", "12345") + assert mlx_distributed_info() == (True, 1, 2) + assert mlx_distributed_uses_mpi() is False + + _clear_mlx_distributed_env(monkeypatch) + ring_hostfile = tmp_path / "ring.json" + ring_hostfile.write_text('[["127.0.0.1:5000"], ["127.0.0.1:5001"]]\n') + monkeypatch.setenv("MLX_RANK", "0") + monkeypatch.setenv("MLX_HOSTFILE", str(ring_hostfile)) + assert mlx_distributed_info() == (True, 0, 2) + assert mlx_distributed_uses_mpi() is False + + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.setenv("MLX_RANK", "1") + monkeypatch.setenv("MLX_IBV_DEVICES", '[["node-a"], ["node-b"]]') + monkeypatch.setenv("MLX_JACCL_COORDINATOR", "node-a:12345") + assert mlx_distributed_info() == (True, 1, 2) + assert mlx_distributed_uses_mpi() is False + + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.setenv("OMPI_COMM_WORLD_RANK", "1") + monkeypatch.setenv("OMPI_COMM_WORLD_SIZE", "2") + assert mlx_distributed_info() == (True, 1, 2) + assert mlx_distributed_uses_mpi() is True + + _clear_mlx_distributed_env(monkeypatch) + monkeypatch.setenv("MLX_RANK", "bad") + monkeypatch.setenv("MLX_WORLD_SIZE", "-3") + assert mlx_distributed_info() == (False, 0, None) + + def test_chat_command_is_registered_with_options(): params = inspect.signature(chatmod.chat).parameters assert "model" in params @@ -751,7 +836,6 @@ def test_chat_server_mode_compare_loads_base_locally(monkeypatch): assert result.exit_code == 0, result.output assert "(compare on)" in result.output - # Only the base model loaded locally, on its own private backend. assert base_loads == [("fake/base", True)] assert streamed == ["base", "tuned"] assert set(closed) == {"http", "base"} @@ -785,6 +869,248 @@ def test_chat_compare_on_mlx_loads_base_model_side_by_side(monkeypatch): assert result.exit_code == 0, result.output assert loads == [("tuned-run", False), ("fake/base", True)] - # Both models answered the turn, via plain generation (no adapter toggle). assert ("base", None) in streamed and ("tuned", None) in streamed assert set(closed) == {"tuned", "base"} + + +@pytest.mark.parametrize( + ("chunk_kind", "expected_exit"), + [ + ("answer", 0), + ("model_text_error", 0), + ("real_error", 1), + ], +) +def test_inference_under_mlx_launch_handles_stream(monkeypatch, chunk_kind, expected_exit): + from unsloth_cli.commands import inference as infermod + from unsloth_cli._inference import ensure_studio_backend_path + + ensure_studio_backend_path() + from core.inference.orchestrator import GenStreamError + + if chunk_kind == "answer": + chunks = ["answer"] + elif chunk_kind == "model_text_error": + # Model output whose visible text starts with "Error:" must not abort. + chunks = ["Error: printed by the model, not a backend failure"] + else: + chunks = [GenStreamError("Error: generation failed")] + + loads, closed = [], [] + + class _FakeBackend: + def stream(self, messages, **kwargs): + return iter(chunks) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setattr( + infermod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr( + infermod, + "load_chat_backend", + lambda model, **kwargs: (loads.append((model, kwargs)), _FakeBackend())[1], + ) + + result = CliRunner().invoke( + _inference_app(), + ["fake-model", "hello", "--tensor-parallel"], + ) + + assert result.exit_code == expected_exit, result.output + assert loads[0][1]["tensor_parallel"] is True + if chunk_kind == "real_error": + assert "generation failed" in result.output + + +def test_chat_under_mlx_launch_nonzero_rank_drains_stdin(monkeypatch): + drains, closed = [], [] + turns = iter( + [ + {"type": "turn", "text": "hi"}, + {"type": "turn", "text": "/exit"}, + ] + ) + + class _FakeChatBackend: + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + assert obj is None + return next(turns) + + def stream(self, messages, **kwargs): + return iter(["hidden"]) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "1") + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr( + chatmod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + monkeypatch.setattr(chatmod, "_drain_available_stdin", lambda: drains.append(True)) + + result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n") + + assert result.exit_code == 0, result.output + assert "Chatting with" not in result.output + assert drains == [True, True] + assert closed == [True] + + +def test_chat_under_mlx_launch_rank0_bypasses_studio_and_prints(monkeypatch): + loads, shares, closed = [], [], [] + + class _FakeChatBackend: + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + shares.append((obj, timeout)) + return obj + + def stream(self, messages, **kwargs): + return iter(["hello"]) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr( + chatmod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr( + chatmod, + "load_chat_backend", + lambda model, **kwargs: (loads.append((model, kwargs)), _FakeChatBackend())[1], + ) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + + result = CliRunner().invoke( + _chat_app(), + ["fake-model", "--tensor-parallel"], + input = "hi\n/exit\n", + ) + + assert result.exit_code == 0, result.output + assert "Chatting with fake-model" in result.output + assert "hello" in result.output + assert loads and loads[0][0] == "fake-model" + assert loads[0][1]["tensor_parallel"] is True + assert shares == [ + ({"type": "turn", "text": "hi"}, None), + ({"type": "turn", "text": "/exit"}, None), + ] + + +@pytest.mark.parametrize( + ("stream_error", "expected_exit"), + [("exception", 1), ("chunk", 1), ("model_text", 0)], +) +def test_chat_under_mlx_launch_exits_on_generation_error(monkeypatch, stream_error, expected_exit): + from unsloth_cli._inference import ensure_studio_backend_path + + ensure_studio_backend_path() + from core.inference.orchestrator import GenStreamError + + closed = [] + + class _FakeChatBackend: + def share_distributed_object( + self, + obj, + *, + timeout = 300.0, + ): + return obj + + def stream(self, messages, **kwargs): + if stream_error == "exception": + raise RuntimeError("generation failed") + if stream_error == "model_text": + # Plain model text starting with "Error:" must not abort the run. + return iter(["Error: printed by the model"]) + return iter([GenStreamError("Error: generation failed")]) + + def close(self): + closed.append(True) + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setattr(chatmod, "resolve_model_config", lambda *a, **k: _FakeConfig()) + monkeypatch.setattr( + chatmod, + "connect_studio_server", + lambda *_a, **_k: (_ for _ in ()).throw(AssertionError("server disabled")), + ) + monkeypatch.setattr(chatmod, "load_chat_backend", lambda *a, **k: _FakeChatBackend()) + monkeypatch.setattr(chatmod, "_compare_needs_second_model", lambda: False) + + result = CliRunner().invoke(_chat_app(), ["fake-model"], input = "hi\n/exit\n") + + assert result.exit_code == expected_exit + if expected_exit: + assert "generation failed" in result.output + assert closed == [True] + + +def test_load_chat_backend_forwards_mlx_distributed_options(monkeypatch): + import unsloth_cli._inference as inference + + calls = [] + + class _FakeBackend: + def load_model(self, **kwargs): + calls.append(kwargs) + return True + + class _FakeModelConfig: + is_gguf = False + + @classmethod + def from_identifier(cls, **_kwargs): + return cls() + + fake_backend = _FakeBackend() + fake_inference = types.ModuleType("core.inference") + fake_inference.get_inference_backend = lambda: fake_backend + fake_utils = types.ModuleType("utils") + fake_utils.__path__ = [] + fake_models = types.ModuleType("utils.models") + fake_models.ModelConfig = _FakeModelConfig + + _set_mlx_nccl_env(monkeypatch, rank = "0") + monkeypatch.setitem(sys.modules, "core", types.ModuleType("core")) + monkeypatch.setitem(sys.modules, "core.inference", fake_inference) + monkeypatch.setitem(sys.modules, "utils", fake_utils) + monkeypatch.setitem(sys.modules, "utils.models", fake_models) + monkeypatch.setattr(inference, "ensure_studio_backend_path", lambda: None) + + inference.load_chat_backend( + "fake-model", + hf_token = None, + max_seq_length = 2048, + load_in_4bit = True, + tensor_parallel = True, + ) + + assert calls[0]["tensor_parallel"] is True + assert calls[0]["mlx_distributed"] is True From 934f879043b289bc6fd699ccfc13f0f363cd655d Mon Sep 17 00:00:00 2001 From: Long Yixing Date: Wed, 8 Jul 2026 18:25:50 +0800 Subject: [PATCH 018/402] feat(mlx): route trainer callbacks (#6929) --- tests/python/test_mlx_public_trainer_api.py | 58 +++++++++++++++++---- unsloth/__init__.py | 50 ++++++++++++++++++ 2 files changed, 99 insertions(+), 9 deletions(-) diff --git a/tests/python/test_mlx_public_trainer_api.py b/tests/python/test_mlx_public_trainer_api.py index 2c33f86af1..89f304c76d 100644 --- a/tests/python/test_mlx_public_trainer_api.py +++ b/tests/python/test_mlx_public_trainer_api.py @@ -115,6 +115,9 @@ def test_mlx_training_arguments_accept_trl_style_kwargs(): def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): """Implemented and falsey inert compatibility kwargs should stay quiet.""" unsloth = _import_mlx_unsloth() + supported_eval_kwargs = {} + if "eval_strategy" in unsloth._MLX_TRAINING_CONFIG_FIELDS: + supported_eval_kwargs = {"eval_strategy": "no", "eval_delay": 1} with warnings.catch_warnings(record = True) as caught: warnings.simplefilter("always") @@ -125,12 +128,16 @@ def test_mlx_training_arguments_do_not_warn_for_implemented_or_falsey_extras(): remove_unused_columns = False, assistant_only_loss = False, completion_only_loss = False, + **supported_eval_kwargs, ) assert args.warmup_steps == 2 assert args.padding_free is False assert args.remove_unused_columns is False assert args.completion_only_loss is False + if supported_eval_kwargs: + assert args.eval_strategy == "no" + assert args.eval_delay == 1 assert caught == [] @@ -705,17 +712,10 @@ def test_mlx_trainer_rejects_unsafe_unsupported_sft_kwargs(): ) -def test_mlx_trainer_rejects_metrics_and_callbacks(): - """Trainer hooks should fail because MLXTrainer cannot honor them yet.""" +def test_mlx_trainer_rejects_compute_metrics(): + """compute_metrics is still unsupported by MLXTrainer.""" unsloth = _import_mlx_unsloth() - with pytest.raises(NotImplementedError, match = "callbacks"): - unsloth.UnslothTrainer( - model = _DummyModel(), - tokenizer = None, - train_dataset = [], - callbacks = [object()], - ) with pytest.raises(NotImplementedError, match = "compute_metrics"): unsloth.UnslothTrainer( model = _DummyModel(), @@ -725,6 +725,46 @@ def test_mlx_trainer_rejects_metrics_and_callbacks(): ) +def test_mlx_trainer_accepts_callbacks(): + """Callbacks are routed to MLXTrainer when the zoo backend supports them.""" + unsloth = _import_mlx_unsloth() + from transformers import TrainerCallback + + if not unsloth._mlx_trainer_supports_kwarg("callbacks"): + pytest.skip("requires unsloth-zoo MLXTrainer callback support") + + class Callback(TrainerCallback): + pass + + trainer = unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [Callback()], + ) + assert any(isinstance(cb, Callback) for cb in trainer.callback_handler.callbacks) + + +def test_mlx_trainer_rejects_callbacks_with_old_zoo(monkeypatch): + """Older unsloth-zoo builds should fail clearly instead of TypeError.""" + unsloth = _import_mlx_unsloth() + from transformers import TrainerCallback + + monkeypatch.setattr( + unsloth, + "_mlx_trainer_supports_kwarg", + lambda name: name != "callbacks", + ) + + with pytest.raises(NotImplementedError, match = "callbacks require"): + unsloth.UnslothTrainer( + model = _DummyModel(), + tokenizer = None, + train_dataset = [], + callbacks = [TrainerCallback()], + ) + + def test_mlx_trainer_rejects_custom_data_collator(): """MLXTrainer owns batching; custom SFT data collators must not be ignored.""" unsloth = _import_mlx_unsloth() diff --git a/unsloth/__init__.py b/unsloth/__init__.py index 04cc600725..de5cd0f61b 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -102,6 +102,7 @@ if _IS_MLX: ) from _e import dataclasses as _dataclasses + import inspect as _inspect import importlib.machinery as _machinery import sys as _sys import types as _types @@ -109,6 +110,30 @@ if _IS_MLX: __version__ = unsloth_zoo.__version__ DEVICE_TYPE = "mlx" + _MLX_TRAINER_ACCEPTS_VAR_KWARGS = False + _MLX_TRAINER_SUPPORTED_KWARGS = frozenset() + try: + _MLX_TRAINER_INIT_PARAMETERS = _inspect.signature(MLXTrainer.__init__).parameters + _MLX_TRAINER_ACCEPTS_VAR_KWARGS = any( + param.kind is _inspect.Parameter.VAR_KEYWORD + for param in _MLX_TRAINER_INIT_PARAMETERS.values() + ) + _MLX_TRAINER_SUPPORTED_KWARGS = frozenset( + name + for name, param in _MLX_TRAINER_INIT_PARAMETERS.items() + if name != "self" + and param.kind + in ( + _inspect.Parameter.POSITIONAL_OR_KEYWORD, + _inspect.Parameter.KEYWORD_ONLY, + ) + ) + except (TypeError, ValueError): + pass + + def _mlx_trainer_supports_kwarg(name): + """Return whether the installed zoo MLXTrainer accepts a kwarg.""" + return _MLX_TRAINER_ACCEPTS_VAR_KWARGS or name in _MLX_TRAINER_SUPPORTED_KWARGS def _is_mlx_cuda_device_target(device): """Return True when a torch .to/.cuda target asks for CUDA on MLX.""" @@ -966,6 +991,7 @@ if _IS_MLX: "args", "formatting_func", "processor", + "callbacks", ) _TRL_SFT_TRAINER_POSITIONAL_KWARGS = ( "model", @@ -985,6 +1011,29 @@ if _IS_MLX: ) _MLX_TRAINER_KWARGS = frozenset(_MLX_TRAINER_POSITIONAL_KWARGS) + def _filter_supported_mlx_trainer_kwargs(trainer_kwargs): + """Drop inert/empty kwargs unsupported by this zoo MLXTrainer.""" + unsupported = { + key: value + for key, value in trainer_kwargs.items() + if not _mlx_trainer_supports_kwarg(key) + } + names = sorted( + key for key, value in unsupported.items() if _is_meaningful_mlx_extra_value(value) + ) + if names: + subject = ", ".join(names) + verb = "requires" if len(names) == 1 else "require" + raise NotImplementedError( + "Unsloth MLX: " + f"{subject} {verb} an unsloth-zoo build with " + "matching MLXTrainer support. Upgrade unsloth-zoo together " + "with unsloth." + ) + for key in unsupported: + trainer_kwargs.pop(key, None) + return trainer_kwargs + def _is_mlx_native_text_collator(collator): """HF pad/copy collators are redundant on MLX; match by class name.""" for klass in type(collator).__mro__: @@ -1162,6 +1211,7 @@ if _IS_MLX: trainer_kwargs, config_kwargs, ignored_kwargs = _split_mlx_trainer_kwargs(kwargs) _raise_unsupported_mlx_trainer_kwargs(ignored_kwargs) + trainer_kwargs = _filter_supported_mlx_trainer_kwargs(trainer_kwargs) trainer_kwargs["args"] = _coerce_mlx_training_args( trainer_kwargs.get("args"), config_kwargs, From 07c8bbbf5a48ee059f2f0e7767f664e24c8b0bf6 Mon Sep 17 00:00:00 2001 From: marcandrelarochelle Date: Wed, 8 Jul 2026 07:05:03 -0400 Subject: [PATCH 019/402] (GRPO) Fix PEFT replacement for TRL >= 1.7.0, add missing compute_aux_loss for TRL >= 1.7.0 (#6904) * Fix PEFT replacement for TRL >= 1.7.0, add missing compute_aux_loss for TRL 1.7.0 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix GRPO for TRL >= 1.7.0: PEFT ref-adapter removal and return arity rl.py: for trl >= 1.7.0, scope the PEFT removal regex to the ref-adapter block only by anchoring the end on ref_param.data.copy_(param.data), so it no longer also deletes the following gradient-checkpointing enable_input_require_grads() block. Neutralize TRL 1.7.0's `if _is_quantized_model:` bf16 cast the same way the existing is_loaded_in_4bit cast is handled. rl_replacements.py: initialize _extra_moe_kwargs before use (it was referenced before assignment whenever compute_aux_loss was passed) and only request output_router_logits when the aux loss is actually wanted. rl_replacements.py: _get_per_token_logps_and_entropies now returns a 3-tuple (logps, entropies, aux_loss) for trl >= 1.7.0 and a 2-tuple for older TRL, matching how every TRL call site unpacks the result. Without this, TRL 1.7.x _generate_and_score_completions unpacks 3 values from a 2-tuple and raises "not enough values to unpack (expected 3, got 2)". * Return zero aux_loss placeholder and drop inference-mode aux collection * GRPO TRL >= 1.7.0: reject router aux-loss opt-in at init; drop zero aux placeholder Unsloth's optimized GRPO forward cannot compute the MoE router auxiliary loss. Previously an explicit opt-in (router_aux_loss_coef > 0) returned a fabricated zero, silently training without the requested load-balancing penalty. Now reject it at trainer init with a clear NotImplementedError, and return None (not zero) for the aux slot of TRL's 3-tuple. Default stays off (coef 0), so the common path is unaffected. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO hidden-states fallback: free ModelOutput before chunked log-softmax The old/ref logprob fallback binds the full ModelOutput (which holds every layer's hidden_states when output_hidden_states=True) and kept it alive across chunked_hidden_states_selective_log_softmax, an avoidable OOM on large models. Extract logits then del outputs in both the text and VLM branches. * Version-compat CI: proactively catch TRL GRPO breakage The existing TRL canary is a static symbol/source grep: it verifies symbols exist but is blind to structural changes (TRL 1.7.0's 2->3-tuple per-token-logps return arity and restructured PEFT ref-adapter block, which the fix in this PR addresses, both slipped past it because the methods still existed). Two additions: - test_trl_grpo_pinned_symbols.py: extend TRL_TAGS to 1.5/1.6/1.7 and pin the exact source-string contracts the rl.py / rl_replacements.py transforms depend on for TRL >= 1.7.0 (PEFT elif ref-adapter block + enable_input_require_grads survival, if _is_quantized_model, aux_loss_enabled anchor, compute_aux_loss arity). A future TRL change fails on main a few days before the PyPI release. - test_trl_grpo_fake_run.py + a version-compat-ci job: fake-CUDA run that drives the real GRPO/SFT/DPO source-transform patchers against latest + main TRL on a CPU-only runner (no training) and asserts the generated Unsloth trainer still satisfies the transform contracts. Catches behavioral regressions the grep cannot see. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fake-run test: use a normal Version import for the aux gate * version-compat CI: fix fake-run job gate + torch-absent collection - Drop the invalid job-level matrix if (matrix is not available in jobs..if -> 'Unrecognized named-value: matrix' fails the whole workflow). Use a single job that runs vs TRL latest always and re-runs vs TRL main only on schedule/dispatch via a step-level github.event_name guard. Validated with actionlint. - Module-level skip the fake-run test when torch is absent so daily-fresh-fetch (pytest-only, collects tests/version_compat/) does not crash on the top-level spoof import. * fake-run test: do not skip on import failure unsloth/trl are installed in the grpo-fake-run job, so a failing import is the import-time drift this canary must catch. Keep only the not-installed find_spec skips; let a real import error fail the test. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO arity gate: regex downgrade + fail loud + CI coverage The TRL < 1.7.0 per-token-logps return downgrade was an exact-string replace anchored on the full return line incl. its comment, so a reformat (e.g. pre-commit) could silently no-op it and ship a 3-tuple to older TRL. Switch to a regex tolerant of comment/whitespace drift, and raise if the anchor stops matching (re.subn count != 1) instead of failing silently. Add a monkeypatched trl_version unit test asserting both arities, since CI only installs TRL >= 1.7.0 and never exercised the downgrade otherwise. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fake-run: give SFT/DPO a real contract, not just ast-parse The SFT/DPO fake patch runs only checked the generated trainer parses. Also assert the shared QLoRA _is_quantized_model bf16 cast is neutralized (TRL 1.7's spelling, present in both sft_trainer and dpo_trainer), so a structural TRL change to that block is caught for SFT/DPO too, not just GRPO. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * GRPO PEFT ref-adapter removal: lower gate to the TRL 1.4.0 floor The elif is_peft_model(model) and args.beta != 0.0: ref-adapter block was introduced in TRL 1.4.0 and is unchanged through 1.7.x, but the removal was gated at >= 1.7.0, so for 1.4 <= TRL < 1.7 the transform fell through to the 0.27 branch (which matches the older if is_peft_available()... form) and silently no-oped: a PEFT + beta != 0 GRPO run then computed the KL reference from the copied ref adapter instead of the base model. Lower the gate to 1.4.0 and keep the 1.7.0-only router aux-loss fail-fast nested. Widen the pinned-symbol contract test to run from 1.4.0 so the covered versions are actually exercised. * [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> Co-authored-by: Daniel Han --- .github/workflows/version-compat-ci.yml | 77 ++++++ .../version_compat/test_trl_grpo_fake_run.py | 246 ++++++++++++++++++ .../test_trl_grpo_pinned_symbols.py | 110 ++++++++ unsloth/models/rl.py | 36 ++- unsloth/models/rl_replacements.py | 42 ++- 5 files changed, 504 insertions(+), 7 deletions(-) create mode 100644 tests/version_compat/test_trl_grpo_fake_run.py diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index e492d21e99..b15d5bfa25 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -285,6 +285,83 @@ jobs: tests/vllm_compat/test_extended_module_imports.py \ -v --tb=short + # Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike + # the static symbol/source greps above, this drives unsloth's actual + # source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only + # runner under the tests/conftest.py spoof harness -- no GPU, no training. + # Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple + # per-token-logps return, restructured PEFT ref-adapter block) by asserting + # the generated Unsloth trainer still satisfies the transform contracts. + grpo-fake-run: + name: GRPO fake-run (latest + main TRL, CPU spoof) + runs-on: ubuntu-latest + timeout-minutes: 18 + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + path: unsloth + - name: Clone unsloth-zoo @ main + run: | + for attempt in 1 2 3; do + rm -rf "$RUNNER_TEMP/unsloth-zoo" + if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \ + "$RUNNER_TEMP/unsloth-zoo"; then + break + fi + if [ "$attempt" -eq 3 ]; then + echo "::error::git clone unsloth-zoo failed after 3 attempts" + exit 1 + fi + delay=$((5 * attempt)) + echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..." + sleep "$delay" + done + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + - name: Install CPU torch + ecosystem + TRL latest + run: | + python -m pip install --upgrade pip + pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \ + 'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10' + # Ecosystem floors unsloth needs; TRL itself is installed last so it + # can pull the transformers/peft it requires. + pip install \ + 'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \ + 'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \ + 'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow + pip install --upgrade trl + pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo" + pip install --no-deps -e ./unsloth + - name: Fake-run vs TRL latest + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + -v --tb=short + # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge + # TRL break does not red every PR. github.event_name is valid in a step if. + - name: Fake-run vs TRL main (scheduled / dispatch only) + if: ${{ github.event_name != 'pull_request' }} + env: + UNSLOTH_IS_PRESENT: '1' + UNSLOTH_COMPILE_DISABLE: '1' + PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python + run: | + pip install --upgrade "git+https://github.com/huggingface/trl" + cd unsloth + python -c "import trl; print('Resolved TRL', trl.__version__)" + PYTHONPATH=. python -m pytest \ + tests/version_compat/test_trl_grpo_fake_run.py \ + -v --tb=short + # Daily-only: same suites but with --strict on importable upstream # tags. Schedule-only so PR jobs stay fast; cron tolerates a flake. daily-fresh-fetch: diff --git a/tests/version_compat/test_trl_grpo_fake_run.py b/tests/version_compat/test_trl_grpo_fake_run.py new file mode 100644 index 0000000000..c79c9955da --- /dev/null +++ b/tests/version_compat/test_trl_grpo_fake_run.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Fake-CUDA GRPO patch run against the *installed* TRL (CPU-only, no training). + +The static symbol/source-string canaries (test_trl_grpo_pinned_symbols.py) +grep raw TRL source; they never execute unsloth's transforms. This test drives +the real pipeline: under the aggressive CUDA spoof it imports unsloth and calls +`_patch_trl_rl_trainers_impl`, which reads the installed GRPOTrainer via +inspect.getsource, applies every rl.py/rl_replacements.py rewrite, and compiles +the result into an UnslothGRPOTrainer. A structural TRL change that slips past +the greps (e.g. TRL 1.7.0's 2->3-tuple return arity, or a restructured PEFT +ref-adapter block) surfaces here as a transform error, a broken generated +source, or a violated contract -- with no GPU and no training run. + +Meant to run in CI against `trl==latest` and `trl @ main` (see +version-compat-ci.yml). The tests/conftest.py harness pre-loads device_type +with DEVICE_COUNT=0 so unsloth's kernel init takes the CPU-safe path. +""" + +from __future__ import annotations + +import ast +import importlib +import importlib.machinery +import importlib.util +import inspect +import sys +import types +from pathlib import Path + +import pytest + + +# daily-fresh-fetch collects tests/version_compat/ with only pytest installed; +# the spoof and the rest of this module need the real torch runtime. Skip the +# whole module cleanly when torch is absent rather than crashing collection. +if importlib.util.find_spec("torch") is None: + pytest.skip("torch not installed; fake-run needs the real runtime", allow_module_level = True) + +# Apply the spoof BEFORE any unsloth-touching import (mirrors +# tests/vllm_compat/test_extended_module_imports.py). +_SPOOF_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_SPOOF_DIR)) +import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402 + +_spoof.apply() + + +def _stub_module(name: str, attrs: dict | None = None) -> None: + if name in sys.modules: + return + m = types.ModuleType(name) + m.__spec__ = importlib.machinery.ModuleSpec(name = name, loader = None, origin = "") + for k, v in (attrs or {}).items(): + setattr(m, k, v) + sys.modules[name] = m + + +_stub_module("torchcodec") + + +def _trl_version(): + import trl + from packaging.version import Version + return Version(trl.__version__.split("+")[0]) + + +def _patch_grpo_and_get_source() -> str: + """Run the GRPO patcher against the installed TRL and return the generated + UnslothGRPOTrainer source. Calls the impl (not the try/except wrapper) so a + transform/compile regression surfaces as a hard error instead of a silent + no-op.""" + import trl.trainer.grpo_trainer as _g + + from unsloth.models import rl as _rl + + _rl._patch_trl_rl_trainers_impl("grpo_trainer") + patched = _g.GRPOTrainer + assert patched.__name__ == "UnslothGRPOTrainer", ( + f"GRPO patch silently no-oped: trl.trainer.grpo_trainer.GRPOTrainer is " + f"{patched.__name__!r}, expected 'UnslothGRPOTrainer' (transform failed " + f"or dispatch key drifted on this TRL)" + ) + # The transformed body (__init__ rewrites, injected per-token-logps) lives in + # the generated module's `_UnslothGRPOTrainer` base + module-level funcs, not + # the thin UnslothGRPOTrainer subclass -- read the whole generated module. + mod = inspect.getmodule(patched) + return inspect.getsource(mod) if mod is not None else inspect.getsource(patched) + + +@pytest.fixture(scope = "module") +def generated_grpo_source(): + if importlib.util.find_spec("unsloth") is None: + pytest.skip("unsloth not installed") + if importlib.util.find_spec("trl") is None: + pytest.skip("trl not installed") + # Do NOT swallow import errors: unsloth is installed here, so a failing + # `import unsloth` is exactly the import-time TRL/transformers drift this + # canary must surface as a failure, not a skip. + import unsloth # noqa: F401 -- _gpu_init bootstrap under spoof + + return _patch_grpo_and_get_source() + + +def test_grpo_patch_generates_valid_source(generated_grpo_source): + """The generated UnslothGRPOTrainer must be syntactically valid Python.""" + ast.parse(generated_grpo_source) + + +def test_grpo_patch_aux_fail_fast_injected(generated_grpo_source): + """TRL >= 1.7.0: rl.py injects a fail-fast for the unsupported MoE router + aux-loss opt-in right after `self.aux_loss_enabled = ...`.""" + from packaging.version import Version + + if _trl_version() < Version("1.7.0"): + pytest.skip("aux_loss_enabled / router_aux_loss_coef are TRL >= 1.7.0") + assert "does not compute the MoE router auxiliary loss" in generated_grpo_source, ( + "aux fail-fast raise missing from generated trainer; rl.py's " + "aux_loss_enabled .replace() anchor did not match this TRL" + ) + + +def test_grpo_patch_three_tuple_return(generated_grpo_source): + """TRL >= 1.7.0 call sites unpack a 3-tuple from + _get_per_token_logps_and_entropies; the injected replacement must return + (logps, entropies, aux_loss).""" + from packaging.version import Version + if _trl_version() >= Version("1.7.0"): + assert "return logprobs.detach(), entropies, aux_loss" in generated_grpo_source, ( + "3-tuple per-token-logps return missing; the arity version-gate in " + "rl_replacements.py did not emit the >=1.7.0 form" + ) + else: + assert ( + "return logprobs.detach(), entropies, aux_loss" not in generated_grpo_source + ), "2-tuple TRL got the 3-tuple return; arity gate mis-fired" + + +def test_grpo_patch_preserves_grad_checkpointing_block(generated_grpo_source): + """The tightened PR #6904 PEFT regex must remove only the ref-adapter init, + not the following enable_input_require_grads gradient-checkpointing block.""" + from packaging.version import Version + + if _trl_version() < Version("1.7.0"): + pytest.skip("ref-adapter elif block is the TRL >= 1.7.0 shape") + assert "enable_input_require_grads" in generated_grpo_source, ( + "gradient-checkpointing enable_input_require_grads() block was swallowed " + "by the PEFT-removal regex (over-reach regression)" + ) + + +def test_grpo_patch_neutralizes_ref_adapter_and_qlora_cast(generated_grpo_source): + """TRL >= 1.7.0: the ref-adapter copy and the hardcoded QLoRA bf16 cast must + both be gone from the generated trainer.""" + from packaging.version import Version + + if _trl_version() < Version("1.7.0"): + pytest.skip("targets the TRL >= 1.7.0 PEFT / _is_quantized_model shapes") + assert ( + "ref_param.data.copy_(param.data)" not in generated_grpo_source + ), "TRL's PEFT ref-adapter init survived; rl.py peft_pattern re.sub no-oped" + assert ( + "if _is_quantized_model:" not in generated_grpo_source + ), "TRL's hardcoded QLoRA bf16 cast survived; rl.py neutralization no-oped" + + +# SFT / DPO: the same source-transform patcher runs on them (a fake patch run, +# no training), so a structural TRL change can break generation. Assert the patch +# produces a valid, importable Unsloth trainer AND that the shared QLoRA +# `_is_quantized_model` bf16 cast is neutralized (TRL 1.7's spelling), which the +# patcher applies to every trainer. Catches "and or others" beyond GRPO. + + +def _patch_and_get_source(trainer_file: str, trainer_cls: str) -> str: + if importlib.util.find_spec("unsloth") is None or importlib.util.find_spec("trl") is None: + pytest.skip("unsloth or trl not installed") + # Let a real import failure fail the test (import-time drift is the target). + import unsloth # noqa: F401 + import trl.trainer # noqa: F401 + + from unsloth.models import rl as _rl + + _rl._patch_trl_rl_trainers_impl(trainer_file) + mod = importlib.import_module(f"trl.trainer.{trainer_file}") + patched = getattr(mod, trainer_cls) + assert patched.__name__ == f"Unsloth{trainer_cls}", ( + f"{trainer_cls} patch silently no-oped on this TRL " + f"(got {patched.__name__!r}); source-transform dispatch drifted" + ) + gen = inspect.getmodule(patched) + src = inspect.getsource(gen) if gen is not None else inspect.getsource(patched) + ast.parse(src) + return src + + +def _assert_quantized_cast_neutralized(src: str, trainer_cls: str) -> None: + from packaging.version import Version + if _trl_version() < Version("1.7.0"): + pytest.skip("pre-1.7.0 spells the QLoRA cast differently (is_loaded_in_4bit)") + assert "if _is_quantized_model:" not in src, ( + f"{trainer_cls}: TRL's hardcoded QLoRA bf16 cast survived; the shared " + f"rl.py `if _is_quantized_model:` -> `if False:` neutralization no-oped" + ) + + +def test_sft_patch_generates_valid_source(): + src = _patch_and_get_source("sft_trainer", "SFTTrainer") + _assert_quantized_cast_neutralized(src, "SFTTrainer") + + +def test_dpo_patch_generates_valid_source(): + src = _patch_and_get_source("dpo_trainer", "DPOTrainer") + _assert_quantized_cast_neutralized(src, "DPOTrainer") + + +# The installed TRL in CI is always >= 1.7.0, so the < 1.7.0 return-arity +# downgrade is never exercised by the fake-run above. Lock both arities by +# monkeypatching rl_replacements.trl_version and re-generating the injected +# _get_per_token_logps_and_entropies source directly (no TRL install needed). +def test_per_token_logps_arity_gate_both_directions(monkeypatch): + if importlib.util.find_spec("unsloth") is None: + pytest.skip("unsloth not installed") + import unsloth # noqa: F401 + from packaging.version import Version + + from unsloth.models import rl_replacements as _rlr + + gate = _rlr.grpo_trainer__get_per_token_logps_and_entropies + + # >= 1.7.0: 3-tuple return kept. + monkeypatch.setattr(_rlr, "trl_version", Version("1.7.0"), raising = False) + src_new = gate("_get_per_token_logps_and_entropies", None) + assert ( + "return logprobs.detach(), entropies, aux_loss" in src_new + ), "3-tuple return missing for TRL >= 1.7.0" + + # < 1.7.0: aux_loss element dropped -> 2-tuple. A no-op downgrade must raise + # (fail loud), never silently ship a 3-tuple to older TRL. + monkeypatch.setattr(_rlr, "trl_version", Version("1.6.0"), raising = False) + src_old = gate("_get_per_token_logps_and_entropies", None) + assert ( + "return logprobs.detach(), entropies # logps, entropies" in src_old + ), "2-tuple return missing for TRL < 1.7.0" + assert ( + "entropies, aux_loss" not in src_old + ), "aux_loss element still present in the TRL < 1.7.0 downgrade" diff --git a/tests/version_compat/test_trl_grpo_pinned_symbols.py b/tests/version_compat/test_trl_grpo_pinned_symbols.py index 834692e2dd..2f935dde49 100644 --- a/tests/version_compat/test_trl_grpo_pinned_symbols.py +++ b/tests/version_compat/test_trl_grpo_pinned_symbols.py @@ -51,10 +51,27 @@ TRL_TAGS = [ "v1.2.0", "v1.3.0", "v1.4.0", + "v1.5.0", + "v1.5.1", + "v1.6.0", + "v1.7.0", # anchor: first release unsloth's TRL>=1.7.0 GRPO patch targets + "v1.7.1", # current PyPI latest "main", ] +def _tag_ge(tag: str, floor: str) -> bool: + """True if `tag` is `main` or a version >= `floor` (e.g. "1.7.0").""" + if tag == "main": + return True + from packaging.version import Version + + try: + return Version(tag.lstrip("v")) >= Version(floor) + except Exception: + return False + + # unsloth/trainer.py + unsloth/models/rl.py rebind these top-level names. @@ -537,3 +554,96 @@ def test_trl_truncate_with_protected_tokens_optional(tag: str): assert src is not None has_it = "truncate_with_protected_tokens" in src _ = has_it # informational; pass either way. + + +# 24-27. TRL >= 1.7.0 GRPO source contracts. Unlike the has_def existence +# checks above, these pin the exact source strings unsloth/models/rl.py and +# rl_replacements.py transform for TRL >= 1.7.0 (the window PR #6904 fixes). +# The 1.7.0 break was invisible to the existence checks because the methods +# still existed -- only their internal structure / return arity changed. If +# TRL restructures one of these, the transform silently no-ops (or the +# generated trainer breaks), so failing here on `main` gives a few-day lead. + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_peft_ref_adapter_block_contract(tag: str): + """rl.py (trl>=1.4.0) strips TRL's PEFT ref-adapter init with a re.DOTALL + regex anchored on `elif is_peft_model(model) and args.beta != 0.0:` ... + `ref_param.data.copy_(param.data)`. Both anchors must exist (else the + regex no-ops and the ref adapter is created under Unsloth), and the + following `enable_input_require_grads` gradient-checkpointing block must + remain present -- the tightened regex must NOT swallow it (PR #6904). The + `elif` block shape appeared in TRL 1.4.0, so this contract runs from there.""" + if not _tag_ge(tag, "1.4.0"): + pytest.skip( + f"{tag}: pre-1.4.0 uses the `if is_peft_available()...` form (rl.py 0.27 branch)" + ) + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert "elif is_peft_model(model) and args.beta != 0.0:" in src, ( + f"{tag}: PEFT ref-adapter `elif` anchor gone; unsloth/models/rl.py " + f"peft_pattern re.sub no-ops and TRL's ref adapter init runs under Unsloth" + ) + assert "ref_param.data.copy_(param.data)" in src, ( + f"{tag}: `ref_param.data.copy_(param.data)` end-anchor gone; " + f"unsloth/models/rl.py peft_pattern loses its DOTALL end match" + ) + assert "enable_input_require_grads" in src, ( + f"{tag}: `enable_input_require_grads` block gone from grpo_trainer.py; " + f"the tightened PR #6904 regex assumed it follows the ref-adapter block" + ) + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_quantized_model_cast_contract(tag: str): + """rl.py (trl>=1.7.0) neutralizes TRL's hardcoded QLoRA bf16 cast + `if _is_quantized_model:` -> `if False:`. A rename leaves the cast active, + which ignores the user's dtype and breaks GradScaler with fp16=True.""" + if not _tag_ge(tag, "1.7.0"): + pytest.skip(f"{tag}: pre-1.7.0 spells the cast differently (is_loaded_in_4bit)") + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert "if _is_quantized_model:" in src, ( + f"{tag}: `if _is_quantized_model:` gone; unsloth/models/rl.py cannot " + f"neutralize TRL's hardcoded QLoRA bf16 cast and it runs under Unsloth" + ) + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_aux_loss_enabled_contract(tag: str): + """rl.py (trl>=1.7.0) appends a fail-fast after + `self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0` so an + explicit MoE router-aux opt-in errors instead of silently training without + the penalty (the optimized forward cannot compute it). A change to this + line drops the guard silently (PR #6904).""" + if not _tag_ge(tag, "1.7.0"): + pytest.skip(f"{tag}: aux_loss_enabled / router_aux_loss_coef added in TRL 1.7.0") + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0" in src, ( + f"{tag}: `aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0` " + f"changed; unsloth/models/rl.py's fail-fast .replace() anchor no-ops" + ) + + +@pytest.mark.parametrize("tag", TRL_TAGS) +def test_trl_grpo_per_token_logps_aux_arity_contract(tag: str): + """TRL 1.7.0 added `compute_aux_loss` to + _get_per_token_logps_and_entropies and made every call site unpack a + 3-tuple. rl_replacements.py version-gates its injected replacement to emit + a 3-tuple for trl>=1.7.0 (2-tuple below). This is the exact change the + has_def existence checks miss: the method still exists, only its arity + changed. If TRL drops/renames the aux return, the gate needs revisiting.""" + if not _tag_ge(tag, "0.20.0"): + pytest.skip(f"{tag}: pre-0.20 uses legacy _get_per_token_logps (2-tuple, no aux)") + src = fetch_text("huggingface/trl", tag, "trl/trainer/grpo_trainer.py") + assert src is not None + assert has_def(src, "_get_per_token_logps_and_entropies", "func"), ( + f"{tag}: _get_per_token_logps_and_entropies missing on TRL >=0.20; " + f"unsloth's per-token-logps injection dispatch key no longer matches" + ) + if _tag_ge(tag, "1.7.0"): + assert "compute_aux_loss" in src, ( + f"{tag}: TRL >=1.7.0 dropped `compute_aux_loss`; the 3-tuple " + f"injection gate in unsloth/models/rl_replacements.py must be revisited" + ) diff --git a/unsloth/models/rl.py b/unsloth/models/rl.py index b5cadf2dea..eeab8fbaca 100644 --- a/unsloth/models/rl.py +++ b/unsloth/models/rl.py @@ -1651,7 +1651,36 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): RLTrainer_source = re.sub(pattern, new_options, RLTrainer_source, flags = re.DOTALL) - if trl_version >= Version("0.27.0"): + if trl_version >= Version("1.4.0"): + # The `elif is_peft_model(model) and args.beta != 0.0:` ref-adapter block + # was introduced in TRL 1.4.0 and is used through 1.7.x. Remove only that + # block, anchored on the final ref_param copy so we do NOT also swallow the + # following gradient-checkpointing enable_input_require_grads() block. + peft_pattern = ( + r"\s*elif is_peft_model\(model\) and args\.beta != 0\.0:" + r".*?" + r"ref_param\.data\.copy_\(param\.data\)" + ) + + replacement_comment = ( + "\n # PEFT initialization logic removed via script for trl >= 1.4.0\n" + ) + + RLTrainer_source = re.sub( + peft_pattern, replacement_comment, RLTrainer_source, flags = re.DOTALL + ) + + if trl_version >= Version("1.7.0"): + # router_aux_loss_coef / aux_loss_enabled were added in TRL 1.7.0. Unsloth's + # optimized GRPO forward cannot compute the MoE router aux loss, so reject + # explicit opt-in (router_aux_loss_coef > 0) at init rather than silently ignoring it. + RLTrainer_source = RLTrainer_source.replace( + "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0", + "self.aux_loss_enabled = is_moe and args.router_aux_loss_coef != 0.0\n" + ' if self.aux_loss_enabled: raise NotImplementedError("Unsloth GRPO does not compute the MoE router auxiliary loss; set router_aux_loss_coef = 0 (the Unsloth default).")', + ) + + elif trl_version >= Version("0.27.0"): peft_pattern = ( r"\s*if is_peft_available\(\) and is_peft_model\(model\) and args\.beta != 0\.0:" r".*?" @@ -1689,6 +1718,11 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"): 'if getattr(model, "is_loaded_in_4bit", False) or getattr(model, "is_loaded_in_8bit", False):', "if False:", ) + # TRL >= 1.7.0 spells the same QLoRA bf16 cast as `if _is_quantized_model:`. + RLTrainer_source = RLTrainer_source.replace( + "if _is_quantized_model:", + "if False:", + ) if RLTrainer_name == "SFTTrainer": original_text = ( diff --git a/unsloth/models/rl_replacements.py b/unsloth/models/rl_replacements.py index 098950de08..ffb845b04f 100644 --- a/unsloth/models/rl_replacements.py +++ b/unsloth/models/rl_replacements.py @@ -1203,6 +1203,8 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1": self._autocast_dtype = torch.float16 + compute_aux_loss = kwargs.get("compute_aux_loss", None) + pixel_values, image_grid_thw = ( kwargs.get("pixel_values", None), kwargs.get("image_grid_thw", None), @@ -1846,7 +1848,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): _extra_vision_kwargs["mm_token_type_ids"] = mm_token_type_ids_chunk with torch.amp.autocast(device_type = "cuda", dtype = self._autocast_dtype): if pixel_values is None: - logits_chunk = unwrapped_model( + outputs = unwrapped_model( input_ids = input_ids_chunk, attention_mask = attention_mask_chunk, pixel_values = pixel_values_chunk, @@ -1854,7 +1856,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): pixel_attention_mask = pixel_attention_mask_chunk, image_sizes = image_sizes_chunk, **_extra_vision_kwargs, - ).logits + ) + + logits_chunk = outputs.logits + del outputs # free hidden_states before chunked log-softmax completion_input_ids_chunk = input_ids_chunk[ :, -(logits_to_keep + max_left_pad) : @@ -1876,7 +1881,7 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): else: # Essentially, for VLMs we do not go via the optimized path in models/, # so we don't encounter the Flash Attn left-padding issue. - logits_chunk = unwrapped_model( + outputs = unwrapped_model( input_ids = input_ids_chunk, attention_mask = attention_mask_chunk, pixel_values = pixel_values_chunk, @@ -1885,7 +1890,10 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): image_sizes = image_sizes_chunk, logits_to_keep = logits_to_keep + 1, **_extra_vision_kwargs, - ).logits + ) + + logits_chunk = outputs.logits + del outputs # free hidden_states before chunked log-softmax logits_chunk = logits_chunk[:, :-1, :] completion_input_ids_chunk = input_ids_chunk[:, -logits_to_keep:] @@ -1914,11 +1922,15 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): all_logprobs_list.append(logprobs_chunk) if logprobs is None: # padded fallback when packing was not used logprobs = torch.cat(all_logprobs_list, dim = 0) + entropies = None os.environ["UNSLOTH_RETURN_HIDDEN_STATES"] = "0" - - return logprobs.detach(), entropies # logps, entropies + # aux loss is unused: it is off by default (router_aux_loss_coef set to 0 in models/rl.py) + # and explicit opt-in is rejected at trainer init, so this is always None (kept in the + # return for TRL >= 1.7.0's 3-tuple contract). + aux_loss = None + return logprobs.detach(), entropies, aux_loss # logps, entropies, aux_loss # input_ids = input_ids[:, -logits_to_keep:] # For transformers<=4.48, logits_to_keep argument isn't supported, so here we drop logits ourselves. # See https://github.com/huggingface/trl/issues/2770 @@ -1937,6 +1949,24 @@ def grpo_trainer__get_per_token_logps_and_entropies(function_name, function): # return logps # compute logprobs for the input tokens function = inspect.getsource(_get_per_token_logps_and_entropies) + if trl_version < Version("1.7.0"): + # TRL < 1.7.0 unpacks (logps, entropies) at every call site; TRL >= 1.7.0 + # always unpacks (logps, entropies, aux_loss). Drop the aux_loss element so + # the return arity matches the installed TRL. Regex tolerates comment / + # whitespace drift on the return line; fail loud if the anchor ever stops + # matching rather than silently shipping a 3-tuple to older TRL. + new_function, n = re.subn( + r"return (logprobs\.detach\(\), entropies), aux_loss[^\n]*", + r"return \1 # logps, entropies", + function, + ) + if n != 1: + raise RuntimeError( + "Unsloth GRPO: could not downgrade the per-token-logps return to a " + f"2-tuple for TRL {trl_version} (matched {n} times, expected 1). The " + "return line changed; update the arity gate in rl_replacements.py." + ) + function = new_function return function From 0e1ed88bb8161d0cb048d46d2b71e50d925eeb46 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 04:06:28 -0700 Subject: [PATCH 020/402] version-compat CI: fake CPU training runs for SFT/GRPO/DPO (#6965) * version-compat CI: fake CPU training runs for SFT/GRPO/DPO Adds a runtime layer on top of the patch-run canary: actually runs trainer.train() for a couple of steps on a CPU-only runner under the CUDA spoof, wrapping a plain tiny HF model in the Unsloth-patched trainer. Exercises the real train() loop (collation, generation, the injected _get_per_token_logps_and_entropies, loss, backward, optimizer) so a TRL or transformers change that breaks the loop at runtime -- not just the source structure -- surfaces here. No GPU, no meaningful numerics. Needs a chain of small CPU shims (eager torch.compile, dynamo suppress, cuda tensor-alloc redirect to CPU, model.for_training/for_inference equivalents) documented inline. Does not exercise Unsloth's Triton/GPU kernels (CPU can't). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cpu fake-train: force adamw_torch + disable dynamo for CPU runner On a real CPU-build torch runner (GitHub CI) two things bit that a CUDA-build torch with GPUs hidden masked locally: - The default optimizer is adamw_8bit (bitsandbytes), whose is_on_gpu() check dies on CPU tensors. Force optim=adamw_torch in all three configs. - import unsloth reinstalls the real torch.compile over the eager passthrough, so the GRPO hot path (chunked_selective_log_softmax) actually compiles and inductor picks the spoofed CUDA device, crashing on device props (gcnArchName). Re-apply the eager passthrough after import and flip torch._dynamo.config.disable so every @torch.compile runs eager at call time. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cpu fake-train: write checkpoints under pytest tmp_path Use pytest's tmp_path for each trainer's output_dir instead of a hardcoded relative temp/ci_* path, so a local pytest run does not leave untracked dirs in the repo tree and the tests are CWD-independent. * version-compat CI: disable dynamo at process level for the fake-run job Set TORCHDYNAMO_DISABLE / TORCH_COMPILE_DISABLE in the fake-run step env so dynamo/inductor is off before conftest.py's early import unsloth, not only via the per-test runtime shim. Defense in depth on the GPU-less runner: the GRPO hot path never compiles regardless of when its functions were decorated. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/workflows/version-compat-ci.yml | 9 + .../version_compat/test_trl_fake_train_cpu.py | 285 ++++++++++++++++++ 2 files changed, 294 insertions(+) create mode 100644 tests/version_compat/test_trl_fake_train_cpu.py diff --git a/.github/workflows/version-compat-ci.yml b/.github/workflows/version-compat-ci.yml index b15d5bfa25..6becccc90a 100644 --- a/.github/workflows/version-compat-ci.yml +++ b/.github/workflows/version-compat-ci.yml @@ -339,12 +339,18 @@ jobs: env: UNSLOTH_IS_PRESENT: '1' UNSLOTH_COMPILE_DISABLE: '1' + # Disable dynamo/inductor at the process level, before conftest.py's early + # `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner + # (defense in depth; the CPU fake-train also flips this at runtime). + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python run: | cd unsloth python -c "import trl; print('Resolved TRL', trl.__version__)" PYTHONPATH=. python -m pytest \ tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ -v --tb=short # `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge # TRL break does not red every PR. github.event_name is valid in a step if. @@ -353,6 +359,8 @@ jobs: env: UNSLOTH_IS_PRESENT: '1' UNSLOTH_COMPILE_DISABLE: '1' + TORCHDYNAMO_DISABLE: '1' + TORCH_COMPILE_DISABLE: '1' PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python run: | pip install --upgrade "git+https://github.com/huggingface/trl" @@ -360,6 +368,7 @@ jobs: python -c "import trl; print('Resolved TRL', trl.__version__)" PYTHONPATH=. python -m pytest \ tests/version_compat/test_trl_grpo_fake_run.py \ + tests/version_compat/test_trl_fake_train_cpu.py \ -v --tb=short # Daily-only: same suites but with --strict on importable upstream diff --git a/tests/version_compat/test_trl_fake_train_cpu.py b/tests/version_compat/test_trl_fake_train_cpu.py new file mode 100644 index 0000000000..4dae696282 --- /dev/null +++ b/tests/version_compat/test_trl_fake_train_cpu.py @@ -0,0 +1,285 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Fake CPU training runs for the Unsloth-patched SFT / GRPO / DPO trainers. + +The patch-run canary (test_trl_grpo_fake_run.py) only compiles + inspects the +generated trainer source. This goes one layer deeper: it actually runs +`trainer.train()` for a couple of steps on a CPU-only runner, under the CUDA +spoof, wrapping a plain (tiny, random-weight) HF model in the Unsloth-patched +trainer. That exercises the real train() loop at runtime -- data collation, +generation (GRPO), the injected `_get_per_token_logps_and_entropies`, loss, +backward, optimizer -- so a TRL or transformers change that breaks the loop +(not just the source structure) surfaces here. No GPU, no meaningful numerics. + +What it does NOT cover: Unsloth's Triton/GPU-optimized model kernels (the +FastLanguageModel fast path) cannot run on CPU, so this validates the +trainer-transform + orchestration layer with a standard forward, not the +optimized kernels. +""" + +from __future__ import annotations + +import os + +# CPU-only: no torch.compile / dynamo (it reaches into the CUDA accelerator), no +# Unsloth kernel compile, no mixed precision. Must be set before torch/unsloth. +os.environ.setdefault("UNSLOTH_COMPILE_DISABLE", "1") +os.environ.setdefault("TORCHDYNAMO_DISABLE", "1") +os.environ.setdefault("TORCH_COMPILE_DISABLE", "1") +os.environ.setdefault("ACCELERATE_MIXED_PRECISION", "no") + +import importlib +import importlib.util +import sys +from pathlib import Path + +import pytest + + +# torch is needed for everything below (daily-fresh-fetch collects this dir with +# only pytest installed); skip the whole module cleanly when it is absent. +if importlib.util.find_spec("torch") is None: + pytest.skip( + "torch not installed; fake CPU train needs the real runtime", allow_module_level = True + ) + +# Apply the CUDA spoof before any unsloth-touching import. +_SPOOF_DIR = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(_SPOOF_DIR)) +import _zoo_aggressive_cuda_spoof as _spoof # noqa: E402 + +_spoof.apply() + +import torch # noqa: E402 + + +# The generated GRPO trainer hard-decorates hot functions with @torch.compile, +# which dynamo processes even under the disable env vars, reaching into +# torch.accelerator (real CUDA) on a GPU-less box. Make torch.compile an eager +# passthrough before unsloth generates/imports the trainer -- same logic, no +# dynamo. (An eager CPU run is exactly what we want here.) +def _eager_compile( + model = None, + *args, + **kwargs, +): + if callable(model): + return model + return lambda fn: fn + + +torch.compile = _eager_compile + +# Belt-and-suspenders: if any @torch.compile still routes through dynamo, let it +# fall back to eager instead of crashing, and stop its stream-capture probe from +# reaching torch.accelerator -> real CUDA on a GPU-less box. +try: + import torch._dynamo # noqa: E402 + torch._dynamo.config.suppress_errors = True +except Exception: + pass +if hasattr(torch, "accelerator"): + torch.accelerator.is_available = lambda *a, **k: False + + +# Redirect any `device="cuda"` tensor allocation / `.to("cuda")` / `.cuda()` to +# CPU. The aggressive spoof deliberately keeps real allocators, but a fake CPU +# train needs cuda-targeted ops (e.g. inductor's init_gpu_context does +# `torch.empty(1, device="cuda")`) to land on CPU instead of erroring. +def _is_cuda_dev(d): + try: + return d is not None and torch.device(d).type == "cuda" + except Exception: + return False + + +for _name in ( + "empty", + "zeros", + "ones", + "full", + "tensor", + "arange", + "randn", + "rand", + "randint", + "empty_like", + "zeros_like", + "ones_like", +): + _orig = getattr(torch, _name, None) + if _orig is None: + continue + + def _redir( + *args, + _orig = _orig, + **kwargs, + ): + if _is_cuda_dev(kwargs.get("device")): + kwargs["device"] = "cpu" + return _orig(*args, **kwargs) + + setattr(torch, _name, _redir) + +_orig_to = torch.Tensor.to + + +def _to_cpu(self, *args, **kwargs): + args = tuple("cpu" if _is_cuda_dev(a) else a for a in args) + if _is_cuda_dev(kwargs.get("device")): + kwargs["device"] = "cpu" + return _orig_to(self, *args, **kwargs) + + +torch.Tensor.to = _to_cpu +torch.Tensor.cuda = lambda self, *a, **k: self + +# Extra CUDA stubs the aggressive spoof lacks, needed to walk a real train(): +# Adam's _cuda_graph_capture_health_check() probes stream capture. +torch.cuda.is_current_stream_capturing = lambda *a, **k: False +try: + import torch.cuda.graphs as _cg # noqa: E402 + _cg._cuda_isCurrentStreamCapturing = lambda *a, **k: False +except Exception: + pass + +# A broken libmlx.so in the shared site-packages crashes transformers' Mac-only +# is_mlx_array probe on Linux; disable it. +try: + import transformers.utils.generic as _g # noqa: E402 + _g._is_mlx_available = False +except Exception: + pass + + +# Dense (non-MoE) tiny model on purpose: MoE models route through Unsloth's +# grouped_gemm Triton kernel, which is CUDA-only and cannot run on a CPU runner. +_MODEL = "hf-internal-testing/tiny-random-LlamaForCausalLM" + + +def _load_plain(): + """Tiny plain HF model + tokenizer on CPU. Skips (not fails) if the model + cannot be fetched -- that is a network/hub issue, not an unsloth regression.""" + from transformers import AutoModelForCausalLM, AutoTokenizer + + try: + tok = AutoTokenizer.from_pretrained(_MODEL) + model = AutoModelForCausalLM.from_pretrained(_MODEL, dtype = torch.float32) + except OSError as e: # hub unreachable / model missing + pytest.skip(f"could not fetch {_MODEL} (network/hub): {str(e)[:150]}") + if tok.pad_token is None: + tok.pad_token = tok.eos_token + # Unsloth's GRPO path calls model.for_training()/for_inference() (added by + # FastLanguageModel). A plain HF model lacks them; supply minimal train/eval + # equivalents so the loop proceeds without the optimized wrapper. + if not hasattr(model, "for_training"): + model.for_training = lambda *a, **k: model.train() + if not hasattr(model, "for_inference"): + model.for_inference = lambda *a, **k: model.eval() + return model.to("cpu"), tok + + +@pytest.fixture(autouse = True) +def _require_stack(): + global torch # the `import torch._dynamo` below would otherwise shadow it as local + if importlib.util.find_spec("unsloth") is None or importlib.util.find_spec("trl") is None: + pytest.skip("unsloth or trl not installed") + # A real import failure is a regression we want to surface, so do not guard it. + import unsloth # noqa: F401 -- patches TRL trainers to the Unsloth variants + + # `import unsloth` reinstalls the real torch.compile (overwriting the eager + # passthrough set at module load), so the GRPO hot path (chunked_selective_ + # log_softmax) would really compile -- and inductor picks the spoofed CUDA + # device, crashing on device props (`gcnArchName`). Re-apply the eager + # passthrough and flip dynamo's call-time kill switch so every @torch.compile + # runs eager regardless of when it was decorated. CPU eager is what we want. + torch.compile = _eager_compile + try: + import torch._dynamo # noqa: E402 + torch._dynamo.config.disable = True + except Exception: + pass + + +def test_sft_trains_on_cpu(tmp_path): + from datasets import Dataset + from trl import SFTConfig, SFTTrainer + + assert SFTTrainer.__name__ == "UnslothSFTTrainer", "SFT patch did not apply" + model, tok = _load_plain() + ds = Dataset.from_list([{"text": "The quick brown fox jumps over the lazy dog."}] * 8) + cfg = SFTConfig( + output_dir = str(tmp_path / "ci_sft"), + per_device_train_batch_size = 2, + max_steps = 2, + logging_steps = 1, + report_to = "none", + save_strategy = "no", + use_cpu = True, + max_length = None, + padding_free = False, + dataset_text_field = "text", + fp16 = False, + bf16 = False, + optim = "adamw_torch", + ) + SFTTrainer(model = model, processing_class = tok, args = cfg, train_dataset = ds).train() + + +def test_grpo_trains_on_cpu(tmp_path): + from datasets import Dataset + from trl import GRPOConfig, GRPOTrainer + + assert GRPOTrainer.__name__ == "UnslothGRPOTrainer", "GRPO patch did not apply" + model, tok = _load_plain() + ds = Dataset.from_list([{"prompt": "hi there"}] * 4) + cfg = GRPOConfig( + output_dir = str(tmp_path / "ci_grpo"), + per_device_train_batch_size = 2, + num_generations = 2, + max_steps = 2, + max_completion_length = 8, + logging_steps = 1, + report_to = "none", + temperature = 1.0, + beta = 0.0, + save_strategy = "no", + use_cpu = True, + use_vllm = False, + fp16 = False, + bf16 = False, + optim = "adamw_torch", + ) + GRPOTrainer( + model = model, + processing_class = tok, + reward_funcs = [lambda completions, **k: [float(len(c)) for c in completions]], + args = cfg, + train_dataset = ds, + ).train() + + +def test_dpo_trains_on_cpu(tmp_path): + from datasets import Dataset + from trl import DPOConfig, DPOTrainer + + assert DPOTrainer.__name__ == "UnslothDPOTrainer", "DPO patch did not apply" + model, tok = _load_plain() + ds = Dataset.from_list( + [{"prompt": "Hi", "chosen": " hello friend", "rejected": " go away"}] * 8 + ) + cfg = DPOConfig( + output_dir = str(tmp_path / "ci_dpo"), + per_device_train_batch_size = 2, + max_steps = 2, + logging_steps = 1, + report_to = "none", + save_strategy = "no", + use_cpu = True, + beta = 0.1, + fp16 = False, + bf16 = False, + optim = "adamw_torch", + ) + DPOTrainer(model = model, processing_class = tok, args = cfg, train_dataset = ds).train() From 6ef09361800a6268ac6ae2f89e37c771e5e516f1 Mon Sep 17 00:00:00 2001 From: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Date: Wed, 8 Jul 2026 12:25:42 +0100 Subject: [PATCH 021/402] Fix OpenClaw start default to local TUI (#6937) * fix: launch OpenClaw local TUI by default * Fix/adjust OpenClaw launch paths for PR #6937 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Default OpenClaw to the local TUI only on a bare invocation The first-arg startswith('-') branch rewrote passthrough globals into a broken command: OpenClaw's grammar is openclaw [--dev] [--profile ] , so 'unsloth start openclaw --profile test' became 'openclaw tui --local --profile test', but tui does not accept --profile (or --dev), so the invocation failed. A leading '--flag value' is ambiguous between a global (--profile test) and a tui option (--message hi), so it cannot be reinterpreted safely. Default to the local TUI only when no passthrough args are given, and forward everything else verbatim so OpenClaw parses it under its own grammar. The bare-launch default (the point of this change) is preserved; explicit subcommands and global flags pass through. --------- Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen Co-authored-by: Wasim Yousef Said --- .github/scripts/agent-guides-drive.sh | 4 ++-- unsloth_cli/commands/start.py | 12 +++++++++++- unsloth_cli/tests/test_start.py | 24 ++++++++++++++++++++++++ 3 files changed, 37 insertions(+), 3 deletions(-) diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index d430d2c172..defdb498c7 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -376,7 +376,7 @@ case "$MODE" in hermes) patch_hermes_tools none invoke_via_connect "$OUT" -z "$PROMPT" ;; openclaw) patch_openclaw_agent notools - invoke_via_connect "$OUT" agent --local --agent ci \ + CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;; *) invoke_via_connect "$OUT" "$PROMPT" ;; esac @@ -449,7 +449,7 @@ case "$MODE" in fi ;; opencode) invoke_via_connect "$out" run "$prompt" ;; hermes) invoke_via_connect "$out" -z "$prompt" ;; - openclaw) invoke_via_connect "$out" agent --local --agent ci \ + openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \ --model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;; *) invoke_via_connect "$out" "$prompt" ;; esac diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 477a47cc3d..764f5c7963 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -1568,7 +1568,17 @@ def openclaw( serve = serve, launch = launch, ) - command = ["openclaw", *ctx.args] + openclaw_args = list(ctx.args) + # Default a bare `unsloth start openclaw` to the local TUI. Anything the caller + # passes through is forwarded verbatim so OpenClaw parses it under its own grammar + # (openclaw [global-flags] [options]): an explicit subcommand, a global + # flag that must precede the command such as --profile/--dev, or a tui option. We + # cannot reinterpret those safely because a leading "--flag value" is ambiguous + # between a global (`--profile test`) and a tui option (`--message hi`); prepending + # `tui --local` would break the global form, so only the empty case is defaulted. + if not openclaw_args: + openclaw_args = ["tui", "--local"] + command = ["openclaw", *openclaw_args] install_hint = ( "iwr -useb https://openclaw.ai/install.ps1 | iex" if os.name == "nt" diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index ee5b442c27..18cb40f18d 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -1634,10 +1634,34 @@ def test_connect_openclaw_no_launch(fake_studio, tmp_path): config = json.loads(config_path.read_text()) assert config["models"]["providers"]["unsloth"]["apiKey"] == "sk-unsloth-feedfacefeedface" assert config["agents"]["defaults"]["model"]["primary"] == f"unsloth/{MODEL['id']}" + assert _launch_command(result.output) == ["openclaw", "tui", "--local"] # OpenAI /v1/chat/completions works on either backend — no GGUF gate. assert not any(c[1].endswith("/api/inference/status") for c in fake_studio) +def test_connect_openclaw_no_launch_keeps_explicit_subcommand(fake_studio): + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch", "crestodian"]) + assert result.exit_code == 0, result.output + assert _launch_command(result.output) == ["openclaw", "crestodian"] + + +def test_connect_openclaw_no_launch_passes_global_flags_through(fake_studio): + # OpenClaw globals (openclaw [--dev] [--profile ] ) precede the + # command, and tui does not accept them, so any passthrough args must be forwarded + # verbatim rather than rewritten into `openclaw tui --local `. + result = CliRunner().invoke(start.start_app, ["openclaw", "--no-launch", "--profile", "test"]) + assert result.exit_code == 0, result.output + assert _launch_command(result.output) == ["openclaw", "--profile", "test"] + + +def test_connect_openclaw_no_launch_keeps_explicit_tui(fake_studio): + result = CliRunner().invoke( + start.start_app, ["openclaw", "--no-launch", "tui", "--message", "hi"] + ) + assert result.exit_code == 0, result.output + assert _launch_command(result.output) == ["openclaw", "tui", "--message", "hi"] + + # ── OpenCode (OpenAI /v1/chat/completions) ─────────────────────────── From e86b7874d433ea1c5a2a49063c07932c36aa63ce Mon Sep 17 00:00:00 2001 From: ErenAta16 Date: Wed, 8 Jul 2026 15:26:50 +0300 Subject: [PATCH 022/402] feat: detect installed coding agent CLIs in Studio settings (#6909) * feat: detect installed coding agent CLIs in Studio settings The API-keys panel only ever showed the "claude" flavor of the `unsloth start` command, so anyone using Codex, OpenCode, OpenClaw, Hermes, or Pi had to manually rewrite the copied command by hand. Add a backend check that looks for each agent's CLI binary on PATH (shutil.which, mirroring the pattern already used elsewhere in studio/backend/utils) and expose it as GET /api/settings/coding-agents. The API-keys panel now renders a picker for all six supported agents, marks the ones it finds installed, and defaults to one of those instead of always falling back to claude. Includes unit tests for the detection helper. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review feedback on coding-agent detection Three fixes from PR review: - detect_installed_coding_agents now treats a PATH lookup failure as "not installed" instead of letting it bubble up and break the settings endpoint; added a regression test for it. - CodingAgentsResponse.agents is now typed as an immutable tuple instead of a list built from one, matching CODING_AGENTS itself. - Fixed a race in the API-keys panel: picking an agent while the installed-CLI check is still in flight could get silently overwritten once that check resolved. A ref now tracks whether the user has made a manual choice, so the auto-detected default only applies before that happens. * Address Codex feedback: GGUF gating and remote-detection scope - codex refuses to launch against a non-GGUF (transformers-backed) model (unsloth_cli's _require_gguf_for_codex), so auto-defaulting to it produced a copy-pasteable command that fails immediately whenever the loaded model isn't GGUF. Add useActiveModelIsGguf() (looks up the active checkpoint in the chat runtime store) and a correction effect that steers the auto-pick away from codex unless the loaded model qualifies, without ever touching a choice the user made by hand. - Detection runs via shutil.which on the Studio backend host, which isn't the same machine as the browser in a tunnel/remote session. Reword the 'installed'/'detected' copy to say so explicitly when the tunnel URL is in use, instead of implying the check ran on the viewer's own device. * Rework auto-default per review: loopback gating + inline GGUF check Replaces the previous approach with the exact shape discussed on the PR: - Export isLoopbackHost/normalizeHost from agent-command.ts. The detection endpoint runs shutil.which on the Studio backend, which only describes the browser's own machine when the base this panel targets resolves to loopback. For a LAN or tunnel/remote base, gate the whole thing off -- don't mark anything as "detected" and don't let it drive the default -- instead of just relabeling the copy. - Drop the separate GGUF-correction effect and useActiveModelIsGguf hook. Read useChatRuntimeStore.getState().activeGgufVariant inline inside the existing detection effect's .then() (so it doesn't need to sit in the effect's deps), and pick the first detected agent that isn't codex unless the loaded model is GGUF, leaving the existing default untouched when no compatible agent is detected. Verified both branches (loopback vs LAN/tunnel base, gguf vs non-gguf, manual pick preserved, no-compatible-agent fallback) with a standalone port of the .then() logic. * Address latest Codex findings: stale detection, model swap, cache - Clear detectedAgents (and skip the network call entirely) when the panel leaves a loopback base, instead of leaving a previous loopback detection result marked 'installed' for a command that now targets a LAN/tunnel/ remote host. - Add a separate, network-free correction effect keyed on the live activeGgufVariant: if codex was auto-picked while a GGUF model was loaded and the user then switches to a transformers-backed model while this panel stays mounted, steer away from codex instead of leaving a command that unsloth_cli's _require_gguf_for_codex will now reject. Never touches a manual pick. - Drop coding-agents.ts's module-lifetime cache. Installed-CLI detection is environment state, not a persisted setting, so a stale positive/negative from before the user installed something (or reopened the tab) is worse than one extra cheap local API call per mount; keep only the in-flight de-dupe for concurrent callers. Verified the correction-effect logic (gguf->non-gguf swap with/without a fallback, still-gguf no-op, manual pick never overridden) with a standalone port of the effect. * Make the codex/GGUF auto-pick symmetric in both directions The correction effect only steered away from codex when the model stopped being GGUF; it never steered back toward codex if the model became GGUF *after* a non-GGUF-gated fallback had already picked something else (e.g. codex is the only detected CLI, a transformers model is loaded so the selection correctly falls back to the claude default, then the user loads a GGUF model while the panel stays mounted -- codex never gets reconsidered). Consolidate into one effect that re-derives the preferred detected agent from scratch whenever detectedAgents or activeGgufVariant changes, in either direction, instead of only reacting to the codex-specific downgrade case. The fetch effect now only populates detectedAgents/availableAgents; this effect is the single source of truth for what gets auto-picked from that list. Never overrides a manual choice. Verified both transition directions plus the manual-pick-survives and initial-detection cases with a standalone port of the derivation logic. * Reset the auto-pick to the default when it stops being trustworthy Two more real gaps from the latest Codex pass on d988f52: - The unified derivation effect only handled the case where a *different* detected agent could take over. If codex was the only detected agent and auto-picked while a GGUF model was loaded, then the model stopped being GGUF, 'preferred' came back undefined and the effect silently left the selection on codex -- exactly the command unsloth_cli's _require_gguf_for_codex now rejects. Fall back to DEFAULT_AGENT in that case instead of leaving it untouched. - Leaving a loopback base cleared detectedAgents (so the 'installed' badges correctly disappear) but left whatever agent had been auto-picked from that now-stale, server-side-only detection still selected. Reset to DEFAULT_AGENT there too, unless the user picked by hand. Introduces a shared DEFAULT_AGENT constant instead of repeating the "claude" literal at each reset site. Verified all five cases (both new resets, both manual-pick-survives variants, and the existing multi-detected-agent fallback still preferring another compatible agent over resetting) with a standalone port of the effects. * Derive GGUF-ness from the actual loaded state, not just the variant string activeGgufVariant only covers an HF-repo GGUF pick (a specific quant variant string). A direct local .gguf file -- custom folder, LM Studio, or drag-drop -- is just as much a GGUF the codex preflight (unsloth_cli's _require_gguf_for_codex) would accept, but it never has a "variant" to report, so it read as non-GGUF here even though /api/inference/status correctly reports is_gguf: true for it. That mismatch could leave a Codex-only install not auto-selected, or reset an auto-picked Codex, for a model that actually supports it. Combined activeGgufVariant with activeNativePathToken (covers the drag-drop/picked-file case) and ggufContextLength (only ever populated when the backend last reported is_gguf: true for the active model, see applyActiveModelStatusToStore) so all three paths a model can be GGUF through are covered, matching the same is_gguf-or-equivalent check hasGgufSource already applies to a staged pick elsewhere in this codebase. * Clear stale native-path token on a non-GGUF status refresh When a native (drag-dropped or picked) GGUF was loaded and the backend later switches to a transformers model outside the UI load path, refresh() adopts the new /api/inference/status via setCheckpoint and applyActiveModelStatusToStore. Those reset activeGgufVariant and ggufContextLength but never clear activeNativePathToken, so the isGguf OR stays true after the switch and a Codex-only detection auto-selects unsloth start codex for a non-GGUF model its preflight rejects. Drop activeNativePathToken in applyActiveModelStatusToStore whenever the status is non-GGUF. A real GGUF load reports is_gguf: true, so its token is preserved (the load path owns it); only a non-GGUF status clears it. * Add the AGPL-3.0 header to the new studio contract test * Fix/adjust agent detection for PR #6909 * [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> Co-authored-by: danielhanchen Co-authored-by: wasimysaid <112766706+wasimysaid@users.noreply.github.com> --- studio/backend/routes/settings.py | 14 ++ studio/backend/tests/test_coding_agents.py | 50 +++++ studio/backend/utils/coding_agents.py | 39 ++++ .../lib/apply-inference-status-to-store.ts | 25 ++- .../features/settings/api/coding-agents.ts | 45 +++++ .../settings/components/agent-command.ts | 4 +- .../settings/components/usage-examples.tsx | 175 +++++++++++++++++- studio/frontend/src/i18n/locales/en.ts | 2 + .../test_chat_response_details_ui_contract.py | 3 + ...usage_examples_agent_detection_contract.py | 21 +++ 10 files changed, 361 insertions(+), 17 deletions(-) create mode 100644 studio/backend/tests/test_coding_agents.py create mode 100644 studio/backend/utils/coding_agents.py create mode 100644 studio/frontend/src/features/settings/api/coding-agents.ts create mode 100644 tests/studio/test_usage_examples_agent_detection_contract.py diff --git a/studio/backend/routes/settings.py b/studio/backend/routes/settings.py index 914699f540..bbee374334 100644 --- a/studio/backend/routes/settings.py +++ b/studio/backend/routes/settings.py @@ -32,6 +32,7 @@ from utils.helper_precache_settings import ( helper_model_disabled_by_env, set_helper_precache_enabled, ) +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents from utils.openai_auto_switch_settings import ( DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, DEFAULT_OPENAI_AUTO_SWITCH_ENABLED, @@ -174,6 +175,19 @@ def update_helper_precache( return _helper_precache_response(enabled) +class CodingAgentsResponse(BaseModel): + # All agents `unsloth start` supports, in the CLI's declared order. + agents: tuple[str, ...] = CODING_AGENTS + # Subset of `agents` whose CLI binary was found on PATH; the frontend uses + # this to default the API-keys panel to a command the user can run as-is. + detected: list[str] + + +@router.get("/coding-agents", response_model = CodingAgentsResponse) +def get_coding_agents(current_subject: str = Depends(get_current_subject)) -> CodingAgentsResponse: + return CodingAgentsResponse(detected = detect_installed_coding_agents()) + + @router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse) def get_openai_auto_switch( current_subject: str = Depends(get_current_subject), diff --git a/studio/backend/tests/test_coding_agents.py b/studio/backend/tests/test_coding_agents.py new file mode 100644 index 0000000000..b19da1dded --- /dev/null +++ b/studio/backend/tests/test_coding_agents.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for coding-agent CLI detection used by the API-keys settings panel.""" + +from unittest.mock import patch + +from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents + + +def test_matches_unsloth_start_subcommands(): + # Each entry must be an actual `unsloth start ` subcommand name + # (unsloth_cli/commands/start.py). Spelled out here rather than imported + # from that module, which pulls in the CLI's heavier dependencies. + assert CODING_AGENTS == ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def test_detects_only_agents_present_on_path(): + installed = {"claude", "opencode"} + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: f"/usr/bin/{name}" if name in installed else None, + ): + assert detect_installed_coding_agents() == ["claude", "opencode"] + + +def test_returns_empty_list_when_nothing_is_installed(): + with patch("utils.coding_agents.shutil.which", return_value = None): + assert detect_installed_coding_agents() == [] + + +def test_preserves_declared_order_regardless_of_path_lookup_order(): + with patch( + "utils.coding_agents.shutil.which", + side_effect = lambda name: name if name in ("pi", "claude", "hermes") else None, + ): + assert detect_installed_coding_agents() == ["claude", "hermes", "pi"] + + +def test_treats_a_path_lookup_error_as_not_installed(): + # An advisory check: shutil.which raising for one entry (e.g. a permission + # error walking a PATH directory) should not take down the whole endpoint, + # and should not stop the remaining agents from being checked. + def flaky_which(name: str): + if name == "codex": + raise OSError("permission denied") + return name if name == "claude" else None + + with patch("utils.coding_agents.shutil.which", side_effect = flaky_which): + assert detect_installed_coding_agents() == ["claude"] diff --git a/studio/backend/utils/coding_agents.py b/studio/backend/utils/coding_agents.py new file mode 100644 index 0000000000..f7dd2f8357 --- /dev/null +++ b/studio/backend/utils/coding_agents.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Detect which `unsloth start ` coding-agent CLIs are on PATH. + +The web UI only ever shows the user the "claude" flavor of the `unsloth start` +command (see agent-command.ts), leaving anyone using Codex, OpenCode, and the +other supported agents to manually edit the copied command. This module gives +the frontend a way to ask which of those CLIs are actually installed so it can +default to one the user can run immediately. +""" + +import shutil + +# Keep in sync with the `unsloth start ` subcommands defined in +# unsloth_cli/commands/start.py. Each entry is the exact executable name that +# subcommand launches, so a hit here means `unsloth start ` can find the +# binary on PATH without the user installing anything first. +CODING_AGENTS: tuple[str, ...] = ("claude", "codex", "openclaw", "opencode", "hermes", "pi") + + +def _is_on_path(agent: str) -> bool: + # shutil.which is documented to return None on a miss, but PATH lookups can + # still raise (e.g. a permission error while probing a directory entry); + # this is an advisory check, so a lookup failure should read as "not + # installed" instead of breaking the settings endpoint. + try: + return shutil.which(agent) is not None + except OSError: + return False + + +def detect_installed_coding_agents() -> list[str]: + """Return the subset of CODING_AGENTS whose CLI binary is on PATH. + + Order follows CODING_AGENTS, not discovery order, so callers can treat the + first entry as the preferred default among the installed agents. + """ + return [agent for agent in CODING_AGENTS if _is_on_path(agent)] diff --git a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts index 9386650fee..60788c23bd 100644 --- a/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts +++ b/studio/frontend/src/features/chat/lib/apply-inference-status-to-store.ts @@ -3,16 +3,19 @@ import { getInferenceStatus } from "../api/chat-api"; import { mergeBackendRecommendedInference } from "../presets/preset-policy"; +import { clampReasoningEffortToLevels } from "../provider-capabilities"; import { CHAT_REASONING_ENABLED_KEY, - loadOptionalBool, type ReasoningEffort, type ReasoningStyle, + loadOptionalBool, resolveToolsEnabledOnLoad, useChatRuntimeStore, } from "../stores/chat-runtime-store"; -import { isMultimodalResponse, type InferenceStatusResponse } from "../types/api"; -import { clampReasoningEffortToLevels } from "../provider-capabilities"; +import { + type InferenceStatusResponse, + isMultimodalResponse, +} from "../types/api"; import type { ChatModelSummary } from "../types/runtime"; type LocalReasoningEffort = Extract; @@ -31,7 +34,10 @@ export function normalizeSpeculativeType( return "ngram"; } if (s === "mtp+ngram") return "mtp+ngram"; - const parts = s.split(",").map((p) => p.trim()).filter(Boolean); + const parts = s + .split(",") + .map((p) => p.trim()) + .filter(Boolean); const hasMtp = parts.some((p) => p === "mtp" || p === "draft-mtp"); const hasNgram = parts.some( (p) => p === "ngram" || p === "ngram-mod" || p === "ngram-simple", @@ -197,6 +203,12 @@ export function applyActiveModelStatusToStore( ggufContextLength: currentGgufContextLength, ggufMaxContextLength, ggufNativeContextLength, + // A non-GGUF status must also drop a stale native-path token: without this the + // isGguf OR (activeGgufVariant || activeNativePathToken || ggufContextLength) + // stays true after switching from a native GGUF to a transformers model, so a + // Codex-only detection would auto-select for a model its preflight rejects. A real + // GGUF load reports is_gguf: true, so its token is preserved (the load path owns it). + ...(status.is_gguf ? {} : { activeNativePathToken: null }), modelRequiresTrustRemoteCode: status.requires_trust_remote_code ?? false, defaultChatTemplate: nextDefaultChatTemplate, loadedIsMultimodal: isMultimodalResponse(status), @@ -245,7 +257,7 @@ export function applyActiveModelStatusToStore( const mid = checkpointId.toLowerCase(); if (mid.includes("qwen3.5") || mid.includes("qwen3.6")) { const sizeMatch = mid.match(/(\d+\.?\d*)\s*b/); - if (sizeMatch && parseFloat(sizeMatch[1]) < 9) { + if (sizeMatch && Number.parseFloat(sizeMatch[1]) < 9) { reasoningDefault = false; } } @@ -281,8 +293,7 @@ export async function tryAdoptServerActiveModel(): Promise { } // Re-check after the await: keep a checkpoint the user picked meanwhile. - const previousCheckpoint = - useChatRuntimeStore.getState().params.checkpoint; + const previousCheckpoint = useChatRuntimeStore.getState().params.checkpoint; if (previousCheckpoint) { return true; } diff --git a/studio/frontend/src/features/settings/api/coding-agents.ts b/studio/frontend/src/features/settings/api/coding-agents.ts new file mode 100644 index 0000000000..ae371b2d3a --- /dev/null +++ b/studio/frontend/src/features/settings/api/coding-agents.ts @@ -0,0 +1,45 @@ +// 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 { authFetch } from "@/features/auth"; +import { readFastApiError } from "@/lib/format-fastapi-error"; + +export type CodingAgentsInfo = { + // Every agent `unsloth start` supports, in the CLI's declared order. + agents: string[]; + // Subset of `agents` whose CLI binary was found on PATH by the backend. + detected: string[]; +}; + +type ApiCodingAgentsInfo = { + agents: string[]; + detected: string[]; +}; + +// Which CLIs are on PATH is environment state, not a persisted setting -- it +// can change any time the user installs something new, so this only +// de-duplicates concurrent in-flight calls (e.g. React strict-mode's double +// mount) rather than caching the result across the module's lifetime. Every +// fresh call (each time a settings panel mounts) re-checks PATH for real. +let inFlightInfo: Promise | null = null; + +function fromApi(info: ApiCodingAgentsInfo): CodingAgentsInfo { + return { agents: info.agents, detected: info.detected }; +} + +async function fetchCodingAgents(): Promise { + const res = await authFetch("/api/settings/coding-agents"); + if (!res.ok) { + throw new Error( + await readFastApiError(res, "Failed to load installed coding agents"), + ); + } + return fromApi(await res.json()); +} + +export async function loadCodingAgents(): Promise { + inFlightInfo ??= fetchCodingAgents().finally(() => { + inFlightInfo = null; + }); + return inFlightInfo; +} diff --git a/studio/frontend/src/features/settings/components/agent-command.ts b/studio/frontend/src/features/settings/components/agent-command.ts index 9e87922970..38b2d73c3b 100644 --- a/studio/frontend/src/features/settings/components/agent-command.ts +++ b/studio/frontend/src/features/settings/components/agent-command.ts @@ -12,7 +12,7 @@ const DEFAULT_AGENT = "claude"; // URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is // "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below. -function normalizeHost(host: string): string { +export function normalizeHost(host: string): string { const lower = host.toLowerCase(); return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower; } @@ -26,7 +26,7 @@ function isDefaultLocalHost(host: string): boolean { } // Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8. -function isLoopbackHost(host: string): boolean { +export function isLoopbackHost(host: string): boolean { if (host === "localhost" || host === "::1") return true; const octets = host.split("."); return ( diff --git a/studio/frontend/src/features/settings/components/usage-examples.tsx b/studio/frontend/src/features/settings/components/usage-examples.tsx index c43dd0f219..b3396c8d9b 100644 --- a/studio/frontend/src/features/settings/components/usage-examples.tsx +++ b/studio/frontend/src/features/settings/components/usage-examples.tsx @@ -16,6 +16,7 @@ import { fetchDeviceType, usePlatformStore } from "@/config/env"; import { useChatRuntimeStore } from "@/features/chat"; import { useT } from "@/i18n"; import type { TranslationKey } from "@/i18n"; +import { isTauri } from "@/lib/api-base"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { Tick02Icon } from "@/lib/tick-icon"; import { cn } from "@/lib/utils"; @@ -25,14 +26,15 @@ import { InformationCircleIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useEffect, useMemo, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { Streamdown } from "streamdown"; +import { loadCodingAgents } from "../api/coding-agents"; import { type OpenAIAutoSwitchSettings, loadOpenAIAutoSwitchSettings, updateOpenAIAutoSwitchSettings, } from "../api/openai-auto-switch"; -import { buildAgentCommand } from "./agent-command"; +import { buildAgentCommand, isLoopbackHost, normalizeHost } from "./agent-command"; type ExampleType = | "curl" @@ -114,6 +116,30 @@ const DOC_LINKS = [ { label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" }, ]; +// Falls back to this list until the backend's installed-CLI check resolves; +// kept in sync with the `unsloth start ` subcommands and with +// CODING_AGENTS in studio/backend/utils/coding_agents.py. +const DEFAULT_AGENTS = [ + "claude", + "codex", + "openclaw", + "opencode", + "hermes", + "pi", +]; +// The agent selection resets to this whenever an auto-pick is no longer +// trustworthy (leaving loopback, or the only compatible detected agent +// stops being compatible) rather than lingering on a stale choice. +const DEFAULT_AGENT = "claude"; +const AGENT_LABELS: Record = { + claude: "Claude Code", + codex: "Codex", + openclaw: "OpenClaw", + opencode: "OpenCode", + hermes: "Hermes", + pi: "Pi", +}; + const j = (s: string): string => JSON.stringify(s); const shSingle = (s: string): string => s.replace(/'/g, "'\\''"); const psSingle = (s: string): string => s.replace(/'/g, "''"); @@ -399,6 +425,17 @@ function useLoadedModelName(): string { }, [checkpoint, ggufVariant]); } +// Backend PATH detection is only safe in the desktop app, where the UI owns +// the local backend. A browser loopback URL may be an SSH/local port forward. +function canUseLocalAgentDetection(base: string): boolean { + if (!isTauri) return false; + try { + return isLoopbackHost(normalizeHost(new URL(base).hostname)); + } catch { + return false; + } +} + const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [ typeof unslothLightTheme, typeof unslothDarkTheme, @@ -443,7 +480,18 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const [copied, setCopied] = useState(false); const [copiedUrl, setCopiedUrl] = useState(false); const [copiedAgent, setCopiedAgent] = useState(false); + const [agent, setAgent] = useState(DEFAULT_AGENT); + const [availableAgents, setAvailableAgents] = + useState(DEFAULT_AGENTS); + const [detectedAgents, setDetectedAgents] = useState([]); + // True once the user has picked an agent themselves; guards the detection + // effect below from clobbering that choice if it resolves afterward. + const agentPickedByUserRef = useRef(false); const [useTunnel, setUseTunnel] = useState(readUseTunnelPref); + const origin = typeof window !== "undefined" ? window.location.origin : ""; + const base = + useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); + const localAgentDetection = canUseLocalAgentDetection(base); // null while loading; the same setting the General tab exposes (shared cache). const [autoSwitch, setAutoSwitch] = useState( null, @@ -454,6 +502,78 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { void fetchDeviceType({ force: true }); }, []); + // Fetching is the only job of this effect: populate availableAgents/ + // detectedAgents (or clear them). Which agent gets auto-picked from that + // list is derived separately below, so it can react to the loaded model + // changing too, not just a fresh fetch. + useEffect(() => { + // Browser loopback URLs can be SSH/local forwards, so only the desktop app + // may use backend PATH checks to mark or auto-pick local agents. + if (!localAgentDetection) { + setDetectedAgents([]); + // A previously auto-picked agent was only ever verified against the + // Studio backend's PATH, which is meaningless now that this panel no + // longer targets a loopback base -- don't leave it selected, but + // never touch a choice the user made by hand. + if (!agentPickedByUserRef.current) { + setAgent(DEFAULT_AGENT); + } + return; + } + + let cancelled = false; + void loadCodingAgents() + .then((info) => { + if (cancelled) return; + setAvailableAgents(info.agents); + setDetectedAgents(info.detected); + }) + .catch(() => { + // Best-effort: keep the default agent list and let the user pick manually. + }); + return () => { + cancelled = true; + }; + }, [localAgentDetection]); + + // Single source of truth for the auto-picked agent, re-derived whenever + // the detected list or the loaded model's GGUF-ness changes -- in either + // direction. `codex` needs a GGUF model (unsloth_cli's + // _require_gguf_for_codex exits otherwise), so it's only preferred once + // the loaded model actually qualifies; loading a GGUF model *after* a + // non-GGUF-gated fallback picked something else re-steers back to codex + // just as loading a non-GGUF model steers away from it. Never overrides a + // choice the user made by hand. + // activeGgufVariant alone only covers an HF-repo GGUF pick (a specific + // quant variant string) -- a direct local .gguf file (custom folder / + // LM Studio / drag-drop) is just as much a GGUF the codex preflight would + // accept, but never has a "variant" to report, and would otherwise read as + // non-GGUF here. activeNativePathToken covers the drag-drop/picked-file + // case; ggufContextLength is only ever populated when the backend's + // /api/inference/status last reported is_gguf: true for the active model + // (see applyActiveModelStatusToStore), so together these three cover every + // path a model can be GGUF through, matching the same is_gguf-or-equivalent + // check hasGgufSource applies to a staged pick. + const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant); + const activeNativePathToken = useChatRuntimeStore((s) => s.activeNativePathToken); + const ggufContextLength = useChatRuntimeStore((s) => s.ggufContextLength); + useEffect(() => { + if (agentPickedByUserRef.current) return; + if (detectedAgents.length === 0) return; + const isGguf = + activeGgufVariant != null || activeNativePathToken != null || ggufContextLength != null; + const preferred = detectedAgents.find((a) => a !== "codex" || isGguf); + if (preferred) { + setAgent(preferred); + } else if (agent === "codex" && !isGguf) { + // codex was auto-picked while a GGUF model was active and it's the + // only detected agent; now that the model isn't GGUF anymore, nothing + // detected is actually runnable, so fall back to the default instead + // of leaving a codex command unsloth_cli will reject. + setAgent(DEFAULT_AGENT); + } + }, [agent, detectedAgents, activeGgufVariant, activeNativePathToken, ggufContextLength]); + useEffect(() => { let cancelled = false; void loadOpenAIAutoSwitchSettings() @@ -470,9 +590,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { const model = useLoadedModelName(); const key = apiKey || KEY_PLACEHOLDER; - const origin = typeof window !== "undefined" ? window.location.origin : ""; - const base = - useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin); const autoSwitchOn = autoSwitch?.enabled ?? false; const snippets = useMemo( @@ -481,8 +598,8 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { ); // Agent command must target the server the panel shows, not the :8888 default. const agentCommand = useMemo( - () => buildAgentCommand(base, key, os), - [base, key, os], + () => buildAgentCommand(base, key, os, agent), + [base, key, os, agent], ); const osAware = OS_AWARE[lang]; @@ -710,6 +827,42 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) { {t("settings.apiKeys.codingAgentsHint")} +
+ {availableAgents.map((id) => { + const installed = detectedAgents.includes(id); + const active = agent === id; + return ( + + ); + })} +
{agentCommand} @@ -727,7 +880,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
- {t("settings.apiKeys.codingAgentsSwap")} + {detectedAgents.length > 0 + ? t("settings.apiKeys.codingAgentsDetectedHint", { + agents: detectedAgents + .map((id) => AGENT_LABELS[id] ?? id) + .join(", "), + }) + : t("settings.apiKeys.codingAgentsSwap")}
diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index b67fd5ca1d..a9b5d839b3 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -445,6 +445,8 @@ export const en = { codingAgentsHint: "Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.", codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.", + codingAgentDetected: "Installed on this machine", + codingAgentsDetectedHint: "Detected on this machine: {agents}.", relativeNever: "never", relativeJustNow: "just now", relativeHoursAgo: "{count}h ago", diff --git a/tests/studio/test_chat_response_details_ui_contract.py b/tests/studio/test_chat_response_details_ui_contract.py index 04301a0de6..89d3000629 100644 --- a/tests/studio/test_chat_response_details_ui_contract.py +++ b/tests/studio/test_chat_response_details_ui_contract.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + """Static contract for the chat response-details action and metadata.""" from __future__ import annotations diff --git a/tests/studio/test_usage_examples_agent_detection_contract.py b/tests/studio/test_usage_examples_agent_detection_contract.py new file mode 100644 index 0000000000..0dc16f0d30 --- /dev/null +++ b/tests/studio/test_usage_examples_agent_detection_contract.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Static contract for API usage-example agent detection scope.""" + +from pathlib import Path + +REPO = Path(__file__).resolve().parents[2] +USAGE_EXAMPLES_TSX = REPO / "studio/frontend/src/features/settings/components/usage-examples.tsx" + + +def test_agent_detection_requires_desktop_scope(): + src = USAGE_EXAMPLES_TSX.read_text(encoding = "utf-8") + assert 'import { isTauri } from "@/lib/api-base"' in src + assert "function canUseLocalAgentDetection(base: string): boolean" in src + helper = src[src.find("function canUseLocalAgentDetection") : src.find("const SHIKI_THEMES")] + assert "if (!isTauri) return false" in helper + assert "isLoopbackHost(normalizeHost(new URL(base).hostname))" in helper + assert "const localAgentDetection = canUseLocalAgentDetection(base)" in src + assert "if (!localAgentDetection)" in src + assert "}, [localAgentDetection]);" in src From 41dd95ea0a588343fc010ee24af4837cc8c08b98 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 05:33:16 -0700 Subject: [PATCH 023/402] Studio: don't pin transformers before the training worker activates the 5.x sidecar (#6968) * Studio: keep transformers off sys.modules until the training worker activates the sidecar The training worker (core/training/worker.py:run_training_process) decides the per-worker Xet env flip during preflight by importing utils/hf_xet_fallback.py, which eagerly imported unsloth_zoo at module load. unsloth_zoo's __init__ imports transformers, so the default transformers 4.57.x was cached in sys.modules before activate_transformers_for_subprocess prepended the 5.x sidecar to sys.path. Since activation only edits sys.path, the already cached module won, and 5.x models failed to load their tokenizer or config: - Qwen3.5 / GLM-4.7 (tokenizer_class TokenizersBackend): "Tokenizer class TokenizersBackend does not exist or is not currently imported." - gemma-4: "... is not supported yet in transformers==4.57.6." Fix: load the shared unsloth_zoo backend lazily (only when a heavy download helper is first used, which is after activation). child_should_disable_xet and the DEFAULT_* constants are defined locally so importing the shim stays light. The download wrappers, the DownloadStallError class, start_watchdog and get_hf_download_state resolve the shared backend on first use, and the degraded no-unsloth_zoo fallback is preserved. Tests: - test_hf_xet_fallback.py: existing suite kept green via the restored _shared_* seam; the GPU-init retry test now triggers the lazy load explicitly; new guard asserts importing child_should_disable_xet does not import transformers/unsloth_zoo. - test_training_worker_import_discipline.py: new invariant test that the worker preflight imports leave transformers unimported, so this class of regression cannot return silently. Runs in studio-backend-ci (CPU only, no network/GPU/weights). * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: CPU-only guard that activation switches transformers to the model's sidecar version Adds test_worker_activates_correct_transformers.py: runs the real worker preflight (from utils.hf_xet_fallback import child_should_disable_xet) plus the real tier detection and activate_transformers_for_subprocess for a transformers-5.x model (Qwen3.5, tier 530), then asserts the in-process transformers actually switched to the 5.x sidecar. A stale pre-activation import leaves 4.57.x pinned and fails the assertion, which is exactly the TokenizersBackend regression (#6951). Self-contained CUDA spoof (mirrors tests/_zoo_aggressive_cuda_spoof.py) forces unsloth_zoo down its full, transformers-importing init path on a GPU-less runner; without it unsloth_zoo degrades and never preloads transformers, masking the bug. A one-line stub sidecar stands in for the 5.x venv, so no GPU, network, weights, or real sidecar are needed. Passes on this fix, fails on buggy main. * Studio: load the repo's canonical CUDA spoof in the correct-version guard Load tests/_zoo_aggressive_cuda_spoof.py (the committed spoof the consolidated CI already relies on) as the single source of truth so the guard matches CI and stays robust on a CPU-only torch wheel, where a partial hand-rolled spoof could miss a torch.cuda call and let the unsloth_zoo import raise (masking the bug). Falls back to a minimal inline spoof for a standalone studio checkout. Verified: passes on this fix, fails on buggy main, and the fallback path passes when the spoof file is absent. * Studio: declare the lazily-resolved xet names so ruff F822 stays green DownloadStallError, start_watchdog and get_hf_download_state are provided via the module __getattr__ (PEP 562), so ruff F822 flagged them as undefined names in __all__ and the Source-lint / pre-commit checks went red. Add annotation-only declarations (no value bound, so __getattr__ still resolves them lazily to the shared unsloth_zoo backend) to mark them defined for the linter while keeping F822 active for the rest of __all__. * Studio: tighten comments on the sidecar-activation fix and its tests * Studio: mirror the new MLX-dispatch preflight import in the import-discipline guard The worker preflight now also runs 'from core.training.training import is_apple_silicon_training_platform, should_use_mlx_training_backend' before it activates the transformers sidecar. Add that import (guarded) to the guard's preflight snippet so the invariant test stays a faithful mirror: a future change that makes core.training.training pull transformers/unsloth_zoo eagerly would then be caught too. Verified clean on the current tree (no leak). --------- Co-authored-by: danielhanchen Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> --- studio/backend/tests/test_hf_xet_fallback.py | 42 ++- .../test_training_worker_import_discipline.py | 81 ++++ ...t_worker_activates_correct_transformers.py | 155 ++++++++ studio/backend/utils/hf_xet_fallback.py | 353 +++++++++++------- 4 files changed, 490 insertions(+), 141 deletions(-) create mode 100644 studio/backend/tests/test_training_worker_import_discipline.py create mode 100644 studio/backend/tests/test_worker_activates_correct_transformers.py diff --git a/studio/backend/tests/test_hf_xet_fallback.py b/studio/backend/tests/test_hf_xet_fallback.py index 4d73213d15..2fff744b64 100644 --- a/studio/backend/tests/test_hf_xet_fallback.py +++ b/studio/backend/tests/test_hf_xet_fallback.py @@ -287,7 +287,9 @@ def test_degrades_when_shared_helper_import_raises_importerror(): def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): """GPU detection in unsloth_zoo's __init__ raises NotImplementedError on a GPU-less host. The shim - retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails.""" + retries under UNSLOTH_ZOO_DISABLE_GPU_INIT=1, restores the env, and degrades if the retry fails. + The backend loads lazily (first use of a heavy helper), so this triggers the load explicitly + before asserting the retry/degrade behavior.""" import importlib import os @@ -321,11 +323,15 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.meta_path.insert(0, finder) try: degraded = importlib.import_module("utils.hf_xet_fallback") - # First attempt without the light env, then a retry with it set. + # Import is light (lazy backend); unsloth_zoo not loaded yet. + assert seen_env == [], seen_env + # First use of a heavy helper triggers the load (attempt without the light env, then a retry + # with it set); accessing DownloadStallError drives it via __getattr__. + stall_error = degraded.DownloadStallError assert seen_env == [None, "1"], seen_env # Both attempts raised -> Studio still boots in degraded mode. - assert issubclass(degraded.DownloadStallError, RuntimeError) - # The env override must not leak past the import. + assert issubclass(stall_error, RuntimeError) + # The env override must not leak past the load. assert os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") is None finally: sys.meta_path.remove(finder) @@ -333,3 +339,31 @@ def test_retries_under_light_gpu_init_when_import_fails(monkeypatch): sys.modules.update(saved) if saved_shim is not None: sys.modules["utils.hf_xet_fallback"] = saved_shim + + +def test_importing_child_should_disable_xet_stays_light(monkeypatch): + """Regression guard for the stale-transformers-sidecar bug: importing the shim (and + ``child_should_disable_xet``) must NOT pull in ``transformers``/``unsloth_zoo``. The worker calls + this at startup to decide the Xet env flip BEFORE activating the sidecar; an eager import here + would cache the default transformers 4.57.x in sys.modules, defeating the sidecar sys.path prepend + and breaking 5.x models (Qwen3.5/GLM/gemma-4).""" + import importlib + + for name in [ + m + for m in list(sys.modules) + if m == "transformers" + or m.startswith("transformers.") + or m == "unsloth_zoo" + or m.startswith("unsloth_zoo.") + or m == "utils.hf_xet_fallback" + ]: + monkeypatch.delitem(sys.modules, name, raising = False) + + mod = importlib.import_module("utils.hf_xet_fallback") + # The lightweight decision works without the heavy backend. + assert mod.child_should_disable_xet({"disable_xet": True}) is True + assert mod.child_should_disable_xet({}) is False + # And nothing heavy was imported as a side effect. + assert "transformers" not in sys.modules, "importing the shim must not import transformers" + assert "unsloth_zoo" not in sys.modules, "importing the shim must not import unsloth_zoo" diff --git a/studio/backend/tests/test_training_worker_import_discipline.py b/studio/backend/tests/test_training_worker_import_discipline.py new file mode 100644 index 0000000000..a047c91704 --- /dev/null +++ b/studio/backend/tests/test_training_worker_import_discipline.py @@ -0,0 +1,81 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: the training worker must not import ``transformers`` before it activates the +transformers sidecar. + +``core/training/worker.py:run_training_process`` runs a preflight (Xet decision, logging, hardware +detection) and only THEN calls ``_activate_transformers_version`` -> ``activate_transformers_for_subprocess``, +which prepends the correct ``.venv_t5_*`` (5.x) sidecar to ``sys.path``. Because activation only edits +``sys.path``, it is a no-op for any module already cached in ``sys.modules``. So if the preflight imports +``transformers`` (directly or transitively via ``unsloth_zoo``), the default 4.57.x gets pinned before +the sidecar is on the path -- and 5.x models (Qwen3.5, GLM-4.7, gemma-4) then fail to load their +tokenizer/config ("Tokenizer class TokenizersBackend does not exist"). + +This regression shipped once when ``utils/hf_xet_fallback.py`` eagerly imported ``unsloth_zoo`` (which +imports ``transformers``) at module load; the worker imports that shim during preflight to decide the +Xet env flip (see issue #6951). This test locks the invariant in a fresh interpreter. It is CPU-only, +needs no network/GPU/weights/sidecars, so it runs in the standard ``studio-backend-ci`` matrix. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend + +# Mirrors run_training_process's imports that run BEFORE _activate_transformers_version (worker.py); +# keep in sync. torch-dependent imports are optional (a no-torch CI shard skips them) but must still +# not drag in transformers. +_PREFLIGHT_SNIPPET = r""" +import sys + +# worker.py: from utils.hf_xet_fallback import child_should_disable_xet (+ call it) +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) + +# worker.py: from loggers.config import LogConfig +from loggers.config import LogConfig # noqa: F401 + +# worker.py: from utils.hardware import hardware (imports torch, not transformers) +try: + from utils.hardware import hardware as _hw # noqa: F401 +except Exception: + pass # torch may be absent in a no-torch shard; the invariant below still applies + +# worker.py: from .training import is_apple_silicon_training_platform, should_use_mlx_training_backend +# (the MLX-dispatch preflight; must also stay clear of transformers). Guarded because it may pull +# unsloth/trl, absent in a minimal shard -- but a partial import that leaked transformers would still +# be caught by the assertion below. +try: + from core.training.training import ( # noqa: F401 + is_apple_silicon_training_platform as _is_apple, + should_use_mlx_training_backend as _use_mlx, + ) +except Exception: + pass + +leaked_tf = sorted(m for m in sys.modules if m == "transformers" or m.startswith("transformers.")) +leaked_zoo = sorted(m for m in sys.modules if m == "unsloth_zoo" or m.startswith("unsloth_zoo.")) +assert not leaked_tf, f"transformers imported during worker preflight (before sidecar activation): {leaked_tf}" +assert not leaked_zoo, f"unsloth_zoo imported during worker preflight (before sidecar activation): {leaked_zoo}" +print("PREFLIGHT_CLEAN") +""" + + +def test_worker_preflight_does_not_import_transformers(): + """A fresh interpreter running the worker's pre-activation imports must leave ``transformers`` + (and ``unsloth_zoo``) unimported, so the 5.x sidecar prepend is not defeated by a stale module.""" + result = subprocess.run( + [sys.executable, "-c", _PREFLIGHT_SNIPPET], + cwd = str(_BACKEND_DIR), + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight imported transformers before sidecar activation.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + assert "PREFLIGHT_CLEAN" in result.stdout, result.stdout diff --git a/studio/backend/tests/test_worker_activates_correct_transformers.py b/studio/backend/tests/test_worker_activates_correct_transformers.py new file mode 100644 index 0000000000..fe7b8dd25a --- /dev/null +++ b/studio/backend/tests/test_worker_activates_correct_transformers.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Invariant: after the training worker runs its preflight and then activates the transformers +sidecar, the in-process ``transformers`` must be the sidecar version the model requires -- not the +default 4.57.x that the base environment ships. + +The CPU-only "does it choose the correct transformers version" guard, stronger than the pure +import-order check in ``test_training_worker_import_discipline.py``: it runs the REAL tier detection +(``get_transformers_tier``) and REAL activation (``activate_transformers_for_subprocess``) for a +transformers-5.x model (Qwen3.5, tier 530) and asserts the version actually switched. It catches the +whole failure family at once: + + * a stale pre-activation ``transformers`` import (the #6951 / ``TokenizersBackend`` regression: an + already-cached 4.57.x defeats the sidecar's ``sys.path`` prepend), + * a wrong tier selected for a 5.x model, and + * activation not actually swapping the resident module. + +Why the CUDA spoof matters (verified): ``unsloth_zoo``'s eager ``import transformers`` only happens on +its full, GPU-present init path. On a GPU-less runner it silently degrades and never preloads +transformers -- which would MASK the stale-import bug (the check would falsely pass). Spoofing +``torch.cuda`` so ``unsloth_zoo`` believes a GPU is present forces the real init path, exposing the +regression on CPU CI. The spoof mirrors ``tests/_zoo_aggressive_cuda_spoof.py`` but is inlined so the +test is self-contained in the ``studio-backend-ci`` matrix (whose conftest does not apply the shared +spoof). No GPU/network/weights/real sidecar needed: a one-line stub sidecar stands in for the 5.x venv, +so we only assert activation lands on it. + +Proven: passes on the fixed tree (active == 5.3.0) and fails on the buggy tree (active == 4.57.x) on +a simulated GPU-less runner. +""" + +from __future__ import annotations + +import subprocess +import sys +from pathlib import Path + +_BACKEND_DIR = Path(__file__).resolve().parent.parent # studio/backend +# Canonical CUDA spoof at the repo root (studio/backend -> studio -> repo root). Loaded by the +# subprocess when present (matches the consolidated CI); absent in a standalone studio checkout, where +# the subprocess falls back to a minimal inline spoof. +_SPOOF_PATH = _BACKEND_DIR.parent.parent / "tests" / "_zoo_aggressive_cuda_spoof.py" + +# Runs in a fresh interpreter with cwd == studio/backend so ``utils.*`` resolves like the worker. +# STUB_HOME (a pytest tmp dir) holds a throwaway ``.venv_t5_530`` sidecar exporting transformers 5.3.0. +_SNIPPET = r""" +import os, sys +sys.path.insert(0, os.getcwd()) + +# CUDA spoof so unsloth_zoo takes its full, transformers-importing init path on a GPU-less runner. +# Without it unsloth_zoo degrades and never preloads transformers, which would MASK the stale-import +# regression under test (verified). Prefer the repo's canonical spoof (single source of truth, and the +# one the consolidated CI already relies on); fall back to a minimal inline spoof so this also works in +# a standalone studio checkout. If torch is absent the fixed tree still passes below; the bug just +# would not be exposable in that shard. +try: + import torch # noqa: F401 + _sp = os.environ.get("SPOOF_PATH") + if _sp and os.path.exists(_sp): + import importlib.util + _spec = importlib.util.spec_from_file_location("_zoo_aggressive_cuda_spoof", _sp) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + _mod.apply() + else: + torch.cuda.is_available = lambda: True + torch.cuda.device_count = lambda: 1 + torch.cuda.current_device = lambda: 0 + torch.cuda.get_device_capability = lambda *a, **k: (8, 0) + torch.cuda.get_device_name = lambda *a, **k: "NVIDIA A100-SPOOFED" + torch.cuda.is_bf16_supported = lambda *a, **k: True + class _Props: + name = "NVIDIA A100-SPOOFED" + major = 8 + minor = 0 + total_memory = 80 * 1024**3 + multi_processor_count = 108 + torch.cuda.get_device_properties = lambda *a, **k: _Props() + torch.cuda.mem_get_info = lambda *a, **k: (0, 80 * 1024**3) +except Exception: + pass +os.environ["UNSLOTH_IS_PRESENT"] = "1" + +# Stub 5.x sidecar: activation only edits sys.path, so a package that merely exports __version__ is +# enough to prove the resident transformers switched to it. +home = os.environ["STUB_HOME"] +pkg = os.path.join(home, ".venv_t5_530", "transformers") +os.makedirs(pkg, exist_ok = True) +with open(os.path.join(pkg, "__init__.py"), "w") as f: + f.write('__version__ = "5.3.0"\n') +os.environ["UNSLOTH_STUDIO_HOME"] = home + +# Faithful worker preflight (worker.py: from utils.hf_xet_fallback import child_should_disable_xet). +# This is the exact stale-import trigger: on the buggy tree it pulls unsloth_zoo -> transformers 4.57.x +# into sys.modules BEFORE activation. +from utils.hf_xet_fallback import child_should_disable_xet +child_should_disable_xet({}) +_tf = sys.modules.get("transformers") +preload = _tf.__version__ if _tf is not None else None + +# Real tier detection + real activation, with the 530 sidecar pointed at the stub above. +import utils.transformers_version as tv +tv._VENV_T5_530_DIR = os.path.join(home, ".venv_t5_530") +tv._ensure_venv_t5_530_exists = lambda: True +tier = tv.get_transformers_tier("Qwen/Qwen3.5-9B", None) +tv.activate_transformers_for_subprocess("Qwen/Qwen3.5-9B", None) + +import transformers +print(f"RESULT tier={tier} preload={preload} active={transformers.__version__}") +""" + + +def _parse(stdout: str) -> dict[str, str]: + for line in stdout.splitlines(): + if line.startswith("RESULT "): + return dict(kv.split("=", 1) for kv in line.split()[1:]) + return {} + + +def test_worker_activates_correct_transformers_version(tmp_path): + """The worker's real preflight + activation for a transformers-5.x model (Qwen3.5, tier 530) must + leave the in-process ``transformers`` on the 5.x sidecar. A stale pre-activation import leaves the + default 4.57.x pinned and fails this assertion -- exactly the #6951 ``TokenizersBackend`` regression.""" + result = subprocess.run( + [sys.executable, "-c", _SNIPPET], + cwd = str(_BACKEND_DIR), + env = { + **__import__("os").environ, + "STUB_HOME": str(tmp_path), + **({"SPOOF_PATH": str(_SPOOF_PATH)} if _SPOOF_PATH.exists() else {}), + }, + capture_output = True, + text = True, + ) + assert result.returncode == 0, ( + "Worker preflight + activation harness crashed.\n" + f"STDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + ) + parsed = _parse(result.stdout) + assert parsed, f"No RESULT line.\nSTDOUT:\n{result.stdout}\nSTDERR:\n{result.stderr}" + + # Correct tier chosen for a transformers-5.x model (pure, deterministic; no network/GPU). + assert parsed["tier"] == "530", ( + f"Wrong transformers tier for Qwen3.5 (expected 530, got {parsed['tier']}). " + "Tier detection regressed." + ) + + # Activation must actually swap the resident transformers to the sidecar version. If a preflight + # import cached 4.57.x first, the sidecar prepend is a no-op and this stays 4.57.x -- the bug. + assert parsed["active"] == "5.3.0", ( + "Sidecar activation did NOT switch the in-process transformers to the model's 5.x version " + f"(active={parsed['active']}, preloaded-before-activation={parsed['preload']}). A pre-activation " + "transformers import (directly or via unsloth_zoo) defeated the sidecar; 5.x models (Qwen3.5, " + "GLM-4.7, gemma-4) then fail with 'Tokenizer class TokenizersBackend does not exist'. See #6951." + ) diff --git a/studio/backend/utils/hf_xet_fallback.py b/studio/backend/utils/hf_xet_fallback.py index 2dd2247396..9bc4a60fad 100644 --- a/studio/backend/utils/hf_xet_fallback.py +++ b/studio/backend/utils/hf_xet_fallback.py @@ -6,6 +6,16 @@ Re-exports the shared API and injects Studio's marker-aware cache purge (``prepare_cache_for_transport``) so the download manager keeps its ``.transport`` marker semantics on the HTTP retry. + +Import discipline: ``unsloth_zoo``'s ``__init__`` eagerly imports ``transformers``. The workers +import this shim at startup (to decide the per-worker Xet env flip) *before* activating the model's +``transformers`` sidecar. Activation only prepends the sidecar to ``sys.path``, so a ``transformers`` +already cached in ``sys.modules`` (via an eager ``unsloth_zoo`` import here) wins -- pinning the +default 4.57.x and regressing Qwen3.5 / GLM-4.7 / gemma-4 training with +``Tokenizer class TokenizersBackend does not exist``. So the shared backend is loaded **lazily** +(``_load_shared``), only on first use of a heavy download helper, i.e. after the sidecar is active. +``child_should_disable_xet`` and the ``DEFAULT_*`` constants are defined locally so importing them +never triggers the heavy load. """ from __future__ import annotations @@ -13,161 +23,230 @@ from __future__ import annotations import threading from typing import Any, Callable, Optional -_shared_import_error = None -try: - import unsloth_zoo.hf_xet_fallback as _shared - _shared_available = True -except Exception as _exc: # noqa: BLE001 - any import failure must degrade, not crash - # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less Studio - # host. The download helper needs none of it, so retry via the light UNSLOTH_ZOO_DISABLE_GPU_INIT - # path before giving up. - _shared_import_error = _exc - import os as _os +# Defaults mirror unsloth_zoo.hf_xet_fallback; plain literals so they resolve (including as +# default args below) without importing unsloth_zoo/transformers. +DEFAULT_GRACE_PERIOD = 10.0 +DEFAULT_HEARTBEAT_INTERVAL = 30.0 +DEFAULT_STALL_TIMEOUT = 180.0 - _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") - _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" - try: - import unsloth_zoo.hf_xet_fallback as _shared - _shared_available = True - _shared_import_error = None - except Exception as _exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF downloads - _shared_import_error = _exc2 - _shared_available = False - finally: - if _prev_gpu_init is None: - _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) - else: - _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init +# --- lazy shared-backend loader ---------------------------------------------------------------- +_shared: Any = None +_shared_available: Optional[bool] = None # None = not yet attempted +_shared_import_error: Optional[BaseException] = None +_load_lock = threading.Lock() -if _shared_available: - # Bind by assignment so each public name shares one module-level binding with the degraded branch. - DEFAULT_GRACE_PERIOD = _shared.DEFAULT_GRACE_PERIOD - DEFAULT_HEARTBEAT_INTERVAL = _shared.DEFAULT_HEARTBEAT_INTERVAL - DEFAULT_STALL_TIMEOUT = _shared.DEFAULT_STALL_TIMEOUT - DownloadStallError = _shared.DownloadStallError - child_should_disable_xet = _shared.child_should_disable_xet - get_hf_download_state = _shared.get_hf_download_state - start_watchdog = _shared.start_watchdog - _shared_hf_hub_download_with_xet_fallback = _shared.hf_hub_download_with_xet_fallback - _shared_snapshot_download_with_xet_fallback = _shared.snapshot_download_with_xet_fallback -else: - # Degrade instead of crashing Studio: plain HF downloads, stall watchdog disabled. Thin stubs, - # not a second copy of the orchestration; recovery returns once unsloth_zoo is upgraded. - import logging as _logging - _logging.getLogger(__name__).warning( - "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " - "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " - "re-enable automatic Xet -> HTTP download recovery.", - _shared_import_error, - ) +def _load_shared() -> bool: + """Import ``unsloth_zoo.hf_xet_fallback`` on demand; return True if available. Deferred so + importing this module at worker startup does not pull transformers in before the sidecar is + activated. Degrades (returns False) rather than crashing when unsloth_zoo is unavailable.""" + global _shared, _shared_available, _shared_import_error + if _shared_available is not None: + return _shared_available + with _load_lock: + if _shared_available is not None: + return _shared_available + try: + import unsloth_zoo.hf_xet_fallback as shared - DEFAULT_HEARTBEAT_INTERVAL = 30.0 - DEFAULT_STALL_TIMEOUT = 180.0 - DEFAULT_GRACE_PERIOD = 10.0 + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc: # noqa: BLE001 - any import failure must degrade, not crash + # unsloth_zoo's __init__ runs torch/GPU detection, which raises on a torch-less/GPU-less + # host. The download helper needs none of it, so retry via UNSLOTH_ZOO_DISABLE_GPU_INIT. + _shared_import_error = exc + import os as _os - class DownloadStallError(RuntimeError): - """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + _prev_gpu_init = _os.environ.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = "1" + try: + import unsloth_zoo.hf_xet_fallback as shared - def child_should_disable_xet(config: dict) -> bool: - return bool(config.get("disable_xet")) + _shared = shared + _shared_available = True + _shared_import_error = None + return True + except Exception as exc2: # noqa: BLE001 - degrade so Studio still boots with plain HF + _shared_import_error = exc2 + _shared_available = False + import logging as _logging - def get_hf_download_state(*args: Any, **kwargs: Any) -> None: - return None # unmeasurable -> the (absent) watchdog never fires + _logging.getLogger(__name__).warning( + "unsloth_zoo.hf_xet_fallback unavailable (%s); the Xet stall watchdog is " + "disabled. Install/upgrade unsloth_zoo (and its torch dependency) to " + "re-enable automatic Xet -> HTTP download recovery.", + _shared_import_error, + ) + return False + finally: + if _prev_gpu_init is None: + _os.environ.pop("UNSLOTH_ZOO_DISABLE_GPU_INIT", None) + else: + _os.environ["UNSLOTH_ZOO_DISABLE_GPU_INIT"] = _prev_gpu_init - def start_watchdog( - *, - on_heartbeat: "Optional[Callable[[str], None]]" = None, - interval: float = DEFAULT_HEARTBEAT_INTERVAL, - xet_disabled: bool = False, - **kwargs: Any, - ) -> "threading.Event": - # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline - # is not tripped during a long download. - stop = threading.Event() - if on_heartbeat is None: - return stop - transport = "https" if xet_disabled else "xet" - def _beat() -> None: - while not stop.wait(interval): - try: - on_heartbeat(f"Downloading ({transport} transport)...") - except Exception: - pass +def child_should_disable_xet(config: dict) -> bool: + """Single source of truth for the per-worker Xet env flip (mirrors + ``unsloth_zoo.hf_xet_fallback.child_should_disable_xet``). Deliberately lightweight: importing or + calling it must NOT pull in unsloth_zoo/transformers, so the worker can decide before activating + the transformers sidecar (see the module docstring).""" + return bool(config.get("disable_xet")) - threading.Thread( - target = _beat, - daemon = True, - name = "hf-xet-degraded-heartbeat", - ).start() + +# --- degraded stubs (used only when unsloth_zoo is unavailable) ------------------------------- +class _DegradedDownloadStallError(RuntimeError): + """Stub mirror so callers' ``except`` clauses resolve; never raised in degraded mode.""" + + +def _degraded_get_hf_download_state(*args: Any, **kwargs: Any) -> None: + return None # unmeasurable -> the (absent) watchdog never fires + + +def _degraded_start_watchdog( + *, + on_heartbeat: "Optional[Callable[[str], None]]" = None, + interval: float = DEFAULT_HEARTBEAT_INTERVAL, + xet_disabled: bool = False, + **kwargs: Any, +) -> "threading.Event": + # No stall detection, but keep emitting heartbeats so the orchestrator's inactivity deadline + # is not tripped during a long download. + stop = threading.Event() + if on_heartbeat is None: return stop + transport = "https" if xet_disabled else "xet" - def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: - return cancel_event is not None and cancel_event.is_set() + def _beat() -> None: + while not stop.wait(interval): + try: + on_heartbeat(f"Downloading ({transport} transport)...") + except Exception: + pass - def _shared_hf_hub_download_with_xet_fallback( - repo_id: str, - filename: str, - token: Optional[str], - *, - repo_type: str = "model", - revision: Optional[str] = None, - cache_dir: Optional[str] = None, - force_download: bool = False, - cancel_event: "Optional[threading.Event]" = None, - **_ignored: Any, - ) -> str: - # Keep the cancellation contract: do not start or return a download once cancelled. - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") + threading.Thread( + target = _beat, + daemon = True, + name = "hf-xet-degraded-heartbeat", + ).start() + return stop - from huggingface_hub import hf_hub_download - path = hf_hub_download( - repo_id = repo_id, - filename = filename, - token = token, - repo_type = repo_type, - revision = revision, - cache_dir = cache_dir, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path +def _degraded_cancelled(cancel_event: "Optional[threading.Event]") -> bool: + return cancel_event is not None and cancel_event.is_set() - def _shared_snapshot_download_with_xet_fallback( - repo_id: str, - *, - revision: Optional[str] = None, - token: Optional[str] = None, - repo_type: str = "model", - cache_dir: Optional[str] = None, - allow_patterns: Optional[Any] = None, - ignore_patterns: Optional[Any] = None, - force_download: bool = False, - cancel_event: "Optional[threading.Event]" = None, - **_ignored: Any, - ) -> str: - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - from huggingface_hub import snapshot_download +def _degraded_hf_hub_download_with_xet_fallback( + repo_id: str, + filename: str, + token: Optional[str], + *, + repo_type: str = "model", + revision: Optional[str] = None, + cache_dir: Optional[str] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + # Keep the cancellation contract: do not start or return a download once cancelled. + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") - path = snapshot_download( - repo_id = repo_id, - repo_type = repo_type, - revision = revision, - token = token, - cache_dir = cache_dir, - allow_patterns = allow_patterns, - ignore_patterns = ignore_patterns, - force_download = force_download, - ) - if _degraded_cancelled(cancel_event): - raise RuntimeError("Cancelled") - return path + from huggingface_hub import hf_hub_download + + path = hf_hub_download( + repo_id = repo_id, + filename = filename, + token = token, + repo_type = repo_type, + revision = revision, + cache_dir = cache_dir, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +def _degraded_snapshot_download_with_xet_fallback( + repo_id: str, + *, + revision: Optional[str] = None, + token: Optional[str] = None, + repo_type: str = "model", + cache_dir: Optional[str] = None, + allow_patterns: Optional[Any] = None, + ignore_patterns: Optional[Any] = None, + force_download: bool = False, + cancel_event: "Optional[threading.Event]" = None, + **_ignored: Any, +) -> str: + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + + from huggingface_hub import snapshot_download + + path = snapshot_download( + repo_id = repo_id, + repo_type = repo_type, + revision = revision, + token = token, + cache_dir = cache_dir, + allow_patterns = allow_patterns, + ignore_patterns = ignore_patterns, + force_download = force_download, + ) + if _degraded_cancelled(cancel_event): + raise RuntimeError("Cancelled") + return path + + +# --- lazy attribute access for the heavy shared API ------------------------------------------- +# ``DownloadStallError`` (class identity matters for ``except``), ``start_watchdog`` and +# ``get_hf_download_state`` come from the shared backend when available, else the degraded stubs. +# Resolved via PEP 562 ``__getattr__`` so ``from utils.hf_xet_fallback import X`` triggers the load +# only for these heavy names, not for ``child_should_disable_xet`` / ``DEFAULT_*``. +_DEGRADED_ATTRS = { + "DownloadStallError": _DegradedDownloadStallError, + "start_watchdog": _degraded_start_watchdog, + "get_hf_download_state": _degraded_get_hf_download_state, +} + +# Annotation-only declarations for the three names above: they bind NO value, so lookup still misses +# and PEP 562 ``__getattr__`` resolves them lazily -- but ruff/pyflakes see them as defined, so listing +# them in ``__all__`` does not trip F822 (while F822 still catches a real typo elsewhere in the list). +DownloadStallError: type +start_watchdog: Any +get_hf_download_state: Any + + +def __getattr__(name: str) -> Any: + if name in _DEGRADED_ATTRS: + if _load_shared(): + return getattr(_shared, name) + return _DEGRADED_ATTRS[name] + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +# Indirection seam the public wrappers call (and tests monkeypatch): lazy-load the shared backend, +# then dispatch to it or the degraded stub. The ``_shared_*`` names preserve the pre-refactor contract. +def _shared_hf_hub_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.hf_hub_download_with_xet_fallback + if _load_shared() + else _degraded_hf_hub_download_with_xet_fallback + ) + return impl(*args, **kwargs) + + +def _shared_snapshot_download_with_xet_fallback(*args: Any, **kwargs: Any) -> str: + impl = ( + _shared.snapshot_download_with_xet_fallback + if _load_shared() + else _degraded_snapshot_download_with_xet_fallback + ) + return impl(*args, **kwargs) __all__ = [ From fcb1152c76417c9ae6d6c649a5036a36110c69c1 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Wed, 8 Jul 2026 09:34:59 -0300 Subject: [PATCH 024/402] Studio: source CPU llama.cpp prebuilts from unslothai/llama.cpp (#6311) * Studio: source CPU llama.cpp prebuilts from the unslothai fork * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: reject unknown Linux CPU arches and keep ROCm-tooling hosts off the CPU prebuilt * Studio: extend the resolve-prebuilt ROCm-tooling guard to Windows * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: let ROCm-SDK-only CPU hosts take the fork CPU prebuilt * Studio: accept windows-arm64 prebuilt kind and refresh stale fork-routing comments * Studio: correct stale fork-routing comments and --resolve-prebuilt help * Refresh stale ggml-org routing comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Daniel Han --- .../workflows/studio-windows-update-smoke.yml | 6 +- .../tests/test_install_resolve_prebuilt.py | 105 +++-------- .../backend/tests/test_llama_cpp_freshness.py | 6 +- studio/backend/tests/test_llama_cpp_update.py | 7 +- studio/install_llama_prebuilt.py | 63 +++---- studio/setup.ps1 | 23 +-- studio/setup.sh | 74 ++------ .../install/test_llama_pr_force_and_source.py | 14 +- tests/studio/install/test_pr4562_bugfixes.py | 18 +- tests/studio/install/test_rocm_support.py | 100 +++++----- tests/studio/install/test_selection_logic.py | 171 +++++++++++++++++- 11 files changed, 322 insertions(+), 265 deletions(-) diff --git a/.github/workflows/studio-windows-update-smoke.yml b/.github/workflows/studio-windows-update-smoke.yml index 888b3d70a3..5b92f1a3e0 100644 --- a/.github/workflows/studio-windows-update-smoke.yml +++ b/.github/workflows/studio-windows-update-smoke.yml @@ -6,9 +6,9 @@ # windows-latest runner: # # 1. install.ps1 --local --no-torch installs Studio AND auto-fetches -# the prebuilt llama.cpp Windows binary (llama-bNNNN-bin-win-cpu- -# x64 from ggml-org/llama.cpp). Hitting the source-build fallback -# is treated as an Unsloth bug -- Studio must always pick the +# the prebuilt llama.cpp Windows binary (app--windows-x64-cpu +# from unslothai/llama.cpp). Hitting the source-build fallback is +# treated as an Unsloth bug -- Studio must always pick the # prebuilt on Windows. # 2. unsloth studio update --local is idempotent. Two consecutive # runs both report "prebuilt up to date and validated", no diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index e9941d9e62..b825172a63 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -1,7 +1,8 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -"""install_llama_prebuilt.py: host->repo mapping and the --resolve-prebuilt mode. +"""install_llama_prebuilt.py: the --resolve-prebuilt probe (plans against the fork +by default; --published-repo overrides). These back the in-app update for source-build (markerless) installs: the backend asks the installer whether an official prebuilt exists for this host without @@ -24,9 +25,7 @@ if str(_studio) not in sys.path: ilp = importlib.import_module("install_llama_prebuilt") -if not hasattr(ilp, "published_repo_for_host") or not hasattr( - ilp, "resolve_simple_install_release_plans" -): +if not hasattr(ilp, "resolve_simple_install_release_plans"): pytest.skip("PR symbols not present - check branch", allow_module_level = True) FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp @@ -56,73 +55,6 @@ def _host(**kw): return ilp.HostInfo(**base) -def test_published_repo_for_host(): - # CPU-only Linux (x64 and arm64) -> ggml-org upstream. - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True)) == UPSTREAM - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_arm64 = True, machine = "aarch64")) - == UPSTREAM - ) - # GPU Linux -> fork. - assert ( - ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_usable_nvidia = True)) - == FORK - ) - assert ilp.published_repo_for_host(_host(is_linux = True, is_x86_64 = True, has_rocm = True)) == FORK - # CPU-only Windows -> ggml-org (setup.ps1: the fork ships no win-cpu bundle). - assert ( - ilp.published_repo_for_host(_host(system = "Windows", is_windows = True, is_x86_64 = True)) - == UPSTREAM - ) - # GPU Windows -> fork. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True, has_usable_nvidia = True) - ) - == FORK - ) - # macOS -> fork regardless of GPU (ggml-org macOS bundles need too-new macOS). - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - # Linux with AMD tooling but no probed GPU -> fork (setup.sh routes on tooling). - assert ( - ilp.published_repo_for_host( - _host(is_linux = True, is_x86_64 = True), linux_amd_tooling_present = True - ) - == FORK - ) - # The tooling hint is Linux-only: Windows CPU stays on ggml-org. - assert ( - ilp.published_repo_for_host( - _host(system = "Windows", is_windows = True, is_x86_64 = True), - linux_amd_tooling_present = True, - ) - == UPSTREAM - ) - - -def test_macos_intel_and_arm_both_route_to_fork(): - # macOS uses the unslothai fork's own Mac prebuilts for BOTH arm64 and Intel; - # there is no longer any upstream-on-macOS default path, so the obsolete - # pre-macOS-26 pin (b9415) is gone. - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64") - ) - == FORK - ) - assert ( - ilp.published_repo_for_host( - _host(system = "Darwin", is_macos = True, is_x86_64 = True, machine = "x86_64") - ) - == FORK - ) - - def test_macos_upstream_pin_only_for_explicit_pre26_upstream(): pre26 = _host( system = "Darwin", @@ -188,15 +120,13 @@ def test_resolve_prebuilt_unavailable(monkeypatch, capsys): assert out["repo"] == FORK -def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): - # CPU-probed Linux host but rocminfo on PATH: the dispatch must route to the - # fork so a HIP source build is not offered an upstream CPU prebuilt. - monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) - monkeypatch.setattr(ilp.shutil, "which", lambda tool: tool == "rocminfo") +def _run_resolve_capture_host(monkeypatch, capsys): + """Drive --resolve-prebuilt and return the host the resolver was handed.""" seen = {} def _resolver(tag, host, repo, published_release_tag): seen["repo"] = repo + seen["host"] = host raise ilp.PrebuiltFallback("no asset") monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver) @@ -207,10 +137,33 @@ def test_resolve_prebuilt_linux_amd_tooling_routes_to_fork(monkeypatch, capsys): ) assert ilp.main() == ilp.EXIT_SUCCESS out = json.loads(capsys.readouterr().out.strip().splitlines()[-1]) + return seen, out + + +def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys): + # CPU-only Linux host (no GPU): the dispatch routes to the fork, which now + # ships the CPU prebuilt -- it no longer falls back to ggml-org upstream. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == FORK assert out["repo"] == FORK +def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, capsys): + # A CPU-only host that merely has ROCm/HIP SDK tools on PATH (no AMD GPU, so + # detect_host leaves has_rocm False) is a valid CPU-prebuilt target. The probe + # must NOT reclassify it as ROCm from tool presence alone and suppress the CPU + # bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source + # build. The host is left CPU-only and resolves against the fork. + monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True)) + monkeypatch.setattr( + ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == FORK + assert seen["host"].has_rocm is False + + # Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer # sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection. diff --git a/studio/backend/tests/test_llama_cpp_freshness.py b/studio/backend/tests/test_llama_cpp_freshness.py index 2a2e113585..08e1334ac9 100644 --- a/studio/backend/tests/test_llama_cpp_freshness.py +++ b/studio/backend/tests/test_llama_cpp_freshness.py @@ -137,9 +137,9 @@ def test_read_install_marker_finds_windows_cmake_layout(tmp_path): @pytest.mark.parametrize("repo", ["unslothai/llama.cpp", "ggml-org/llama.cpp"]) def test_read_install_marker_carries_published_repo_dynamically(tmp_path, repo): - # The freshness check queries whichever release repo the marker records, - # so CUDA (unslothai), CPU/macOS (ggml-org), and ROCm all get the right - # "latest" tag. + # The freshness check queries whichever release repo the marker records: + # new installs record the fork, legacy CPU/macOS markers still say ggml-org, + # and both must get the right "latest" tag. install_dir = tmp_path / "llama.cpp" _write_marker(install_dir, tag = "b9000", published_repo = repo) bin_path = _fake_binary(install_dir, layout = "cmake") diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 4ceffbf75b..5138e90471 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -594,9 +594,10 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path): def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path): - # CPU installs come from ggml-org. Re-running into the same install-dir/repo - # reproduces the same CPU bundle; --cpu-fallback (which force-drops GPU - # detection) is reserved for setup.sh's arm64 rescue and must not appear here. + # Legacy CPU installs recorded a ggml-org marker (new installs use the fork). + # Re-running into the same install-dir/repo reproduces the same CPU bundle; + # --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's + # arm64 rescue and must not appear here. cmd = _capture_install_cmd( monkeypatch, tmp_path, diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index e40cb3083e..6c75e6c394 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -165,9 +165,9 @@ def env_int( # errors. Only use "master" temporarily when the latest release is missing # support for a new model architecture. DEFAULT_LLAMA_TAG = os.environ.get("UNSLOTH_LLAMA_TAG", "latest") -# Default published repo for prebuilt release resolution. Linux uses -# Unsloth prebuilts; setup.sh/setup.ps1 pass --published-repo explicitly -# for macOS/Windows to override with ggml-org/llama.cpp when needed. +# Default published repo for prebuilt release resolution. Every host plans +# its prebuilt against the Unsloth fork; setup.sh/setup.ps1 pass it via +# --published-repo. ggml-org is reachable only via an explicit override. DEFAULT_PUBLISHED_REPO = "unslothai/llama.cpp" DEFAULT_PUBLISHED_TAG = os.environ.get("UNSLOTH_LLAMA_RELEASE_TAG") DEFAULT_PUBLISHED_MANIFEST_ASSET = os.environ.get( @@ -3135,21 +3135,6 @@ def _apply_host_overrides( return host -def published_repo_for_host(host: HostInfo, *, linux_amd_tooling_present: bool = False) -> str: - """The release repo setup.sh / setup.ps1 pick for this host: macOS always the - fork (ggml-org macOS bundles need too-new macOS); else CPU-only Linux/Windows - -> ggml-org upstream (the fork ships no CPU bundle) and any usable GPU (NVIDIA - or ROCm) -> the fork. linux_amd_tooling_present mirrors setup.sh routing Linux - hosts that expose AMD tooling (rocminfo/amd-smi/hipconfig/hipinfo) to the fork - even when the probe cannot confirm an active GPU. Mirrors the shell routing.""" - if host.is_macos: - return DEFAULT_PUBLISHED_REPO - has_gpu = ( - host.has_usable_nvidia or host.has_rocm or (host.is_linux and linux_amd_tooling_present) - ) - return DEFAULT_PUBLISHED_REPO if has_gpu else UPSTREAM_REPO - - def pick_windows_cuda_runtime(host: HostInfo) -> str | None: if not host.driver_cuda_version: return None @@ -4015,6 +4000,9 @@ def resolve_release_asset_choice( published_choice = published_rocm_choice_for_host(release, host, "windows-rocm") else: published_choice = published_asset_choice_for_kind(release, "windows-cpu") + elif host.is_windows and host.is_arm64: + # Windows arm64 has no GPU prebuilt, so it always takes the CPU bundle. + published_choice = published_asset_choice_for_kind(release, "windows-arm64") elif host.is_macos and host.is_arm64: published_choice = published_asset_choice_for_kind(release, "macos-arm64") elif host.is_macos and host.is_x86_64: @@ -6127,8 +6115,13 @@ def _linux_published_attempts(host: HostInfo, bundle: PublishedReleaseBundle) -> # CPU-only host. A usable-NVIDIA host never reaches here -- if its CUDA # selection produced nothing we want an empty attempt list so the caller # source-builds with CUDA, not a CPU-only binary silently installed on a - # GPU host (mirrors the ROCm branch, and Windows NVIDIA). - cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu") + # GPU host (mirrors the ROCm branch, and Windows NVIDIA). Only x86_64 and + # arm64 have a CPU bundle; any other Linux arch (ppc64le, riscv64, s390x) + # has none, so leave attempts empty and source-build rather than hand it + # the x86_64 linux-cpu binary (the Linux preflight checks libraries, not + # ELF arch, so a wrong-arch binary would not be caught). + kind = "linux-cpu" if host.is_x86_64 else "linux-arm64" if host.is_arm64 else None + cpu_choice = published_asset_choice_for_kind(bundle, kind) if kind else None if cpu_choice is not None: attempts.append(cpu_choice) return attempts @@ -6143,9 +6136,9 @@ def _fork_manifest_release_plans( max_release_fallbacks: int = DEFAULT_MAX_PREBUILT_RELEASE_FALLBACKS, ) -> tuple[str, list[InstallReleasePlan]]: """Manifest-reading branch of resolve_simple_install_release_plans, used for - the fork's bundles whose GPU/arch coverage lives in - llama-prebuilt-manifest.json rather than in the filename: arm64 CUDA, Windows - CUDA, per-gfx ROCm, and macOS. Linux x64 takes the faster filename path.""" + every fork host: all of the fork's bundles describe their GPU/arch coverage + in llama-prebuilt-manifest.json rather than in the asset filename (CPU, + x64/arm64 CUDA, Windows CUDA, per-gfx ROCm, and macOS).""" requested_tag = normalized_requested_llama_tag(llama_tag) allow_older_release_fallback = requested_tag == "latest" and not published_release_tag release_limit = max(1, max_release_fallbacks) @@ -6714,8 +6707,8 @@ def install_prebuilt( log( f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install" ) - # Single resolver: linux-x64 takes the fast filename path internally, - # every other fork host reads the manifest. + # Single resolver: every fork host selects from the release manifest; + # an explicit ggml-org override selects by asset filename instead. requested_tag, release_plans = resolve_simple_install_release_plans( llama_tag, host, @@ -6903,8 +6896,8 @@ def parse_args() -> argparse.Namespace: const = "latest", help = ( "Report whether an official prebuilt exists for this host without " - "downloading. Picks the host's published repo when --published-repo " - "is left at the default. Use --output-format json." + "downloading. Plans against --published-repo (defaults to the " + "fork). Use --output-format json." ), ) parser.add_argument( @@ -6992,24 +6985,16 @@ def main() -> int: return EXIT_SUCCESS if args.resolve_prebuilt is not None: - # Host-aware "is a prebuilt available" probe, no download. A default repo - # means "pick the repo for this host"; PrebuiltFallback == source build. + # Host-aware "is a prebuilt available" probe, no download. Every host now + # plans against the fork (args.published_repo defaults to it); an explicit + # --published-repo overrides. PrebuiltFallback == source build. host = _apply_host_overrides( detect_host(), override_has_rocm = args.has_rocm, override_rocm_gfx = args.rocm_gfx, force_cpu = args.cpu_fallback, ) - # setup.sh routes Linux hosts with AMD tooling to the fork even when no GPU - # is probed; mirror that so a HIP source build is not offered a CPU prebuilt. - amd_tooling = host.is_linux and any( - shutil.which(t) for t in ("rocminfo", "amd-smi", "hipconfig", "hipinfo") - ) - repo = ( - published_repo_for_host(host, linux_amd_tooling_present = amd_tooling) - if args.published_repo == DEFAULT_PUBLISHED_REPO - else args.published_repo - ) + repo = args.published_repo try: _requested, plans = resolve_simple_install_release_plans( args.resolve_prebuilt, host, repo, args.published_release_tag or "" diff --git a/studio/setup.ps1 b/studio/setup.ps1 index bb1e88cc4e..07dcb17335 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -3088,12 +3088,11 @@ $LlamaCppDir = Join-Path $UnslothHome "llama.cpp" $NeedLlamaSourceBuild = $false $SkipPrebuiltInstall = $false $RequestedLlamaTag = if ($env:UNSLOTH_LLAMA_TAG) { $env:UNSLOTH_LLAMA_TAG } else { $DefaultLlamaTag } -# GPU Windows (CUDA / ROCm) installs the fork's app-* prebuilts; CPU-only stays -# on ggml-org (the fork ships no windows-cpu bundle). Mirrors setup.sh's routing. -# A resolved gfx arch counts as a GPU host even when $HasROCm is false (Adrenalin -# driver only, no HIP runtime): the fork's per-gfx bundle ships its own runtime, -# so route there instead of ggml-org / a CPU build. -$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) { "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" } +# Every host installs the fork's app-* prebuilts now: GPU Windows (CUDA / ROCm) +# already did, and the fork now also ships the CPU bundles for Windows x64 and +# arm64 (windows-cpu / windows-arm64). ggml-org artifacts are no longer used by +# default. Mirrors setup.sh's routing. +$HelperReleaseRepo = "unslothai/llama.cpp" $LlamaPr = if ($env:UNSLOTH_LLAMA_PR) { $env:UNSLOTH_LLAMA_PR.Trim() } else { "" } $LlamaPrForce = if ($env:UNSLOTH_LLAMA_PR_FORCE) { $env:UNSLOTH_LLAMA_PR_FORCE.Trim() } else { $DefaultLlamaPrForce } @@ -3283,11 +3282,13 @@ if ($LocalLlamaCppLinked) { # treat a valid ROCm install as mismatched. A name-inferred gfx # arch (Adrenalin-only, no confirmed runtime) still counts as # ROCm-capable -- the ROCm prebuilt bundles its own runtime, - # mirroring the --rocm-gfx forward below. NOTE: this block is - # currently inert -- write_prebuilt_metadata does not persist an - # install_kind key, so $existingKind is always null. If that changes, - # add the remaining host kinds (e.g. windows-arm64) before relying on it. - $expectedKinds = if ($HasROCm -or $script:ROCmGfxArch) { @("windows-rocm", "windows-hip") } elseif ($HasNvidiaSmi) { @("windows-cuda") } else { @("windows-cpu") } + # mirroring the --rocm-gfx forward below. The CPU branch covers both + # the x64 windows-cpu and arm64 windows-arm64 bundles (Windows arm64 + # has no GPU prebuilt). NOTE: this block is currently inert -- + # write_prebuilt_metadata does not persist an install_kind key, so + # $existingKind is always null; keep $expectedKinds in sync with the + # kinds install_llama_prebuilt.py installs before relying on it. + $expectedKinds = if ($HasROCm -or $script:ROCmGfxArch) { @("windows-rocm", "windows-hip") } elseif ($HasNvidiaSmi) { @("windows-cuda") } else { @("windows-cpu", "windows-arm64") } if ($existingKind -and ($existingKind -notin $expectedKinds)) { substep "Removing mismatched llama.cpp install (found '$existingKind', need one of: $($expectedKinds -join ', '))..." Remove-Item -Recurse -Force -LiteralPath $LlamaCppDir -ErrorAction SilentlyContinue diff --git a/studio/setup.sh b/studio/setup.sh index 6a74cd2296..d244e3cdcf 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1176,63 +1176,19 @@ _REQUESTED_LLAMA_TAG="${UNSLOTH_LLAMA_TAG:-${_DEFAULT_LLAMA_TAG}}" _HOST_SYSTEM="$(uname -s 2>/dev/null || true)" _HOST_MACHINE="$(uname -m 2>/dev/null || true)" -# Pick the release repo install_llama_prebuilt.py plans against. -# The fork ships CUDA (Linux x64/arm64, Windows), ROCm (Linux/Windows) and -# macOS bundles. Only the plain CPU/Vulkan bundles still come from ggml-org, so -# CPU-only Linux (x86_64 and arm64) routes there; GPU Linux, Windows and macOS -# use unslothai. -_LINUX_HAS_GPU=false -# Route to the fork only for a usable GPU. NVIDIA counts only when a device is -# actually enumerated and not hidden via CUDA_VISIBLE_DEVICES=""/-1 -# (_setup_nvidia_usable, from _setup_has_usable_nvidia_gpu above) -- mirroring -# install_llama_prebuilt.py's has_usable_nvidia. Mere nvidia-smi presence -# (CPU-only CUDA-toolkit containers, broken drivers) or a hidden GPU therefore -# takes the ggml-org CPU prebuilt instead of a slow source build. AMD is -# deliberately left on tooling presence, not usability: an unusable NVIDIA host -# has a good CPU prebuilt to fall back to, whereas tightening AMD would regress -# ROCm hosts exposing only hipconfig/hipinfo into an unnecessary CPU build. -if [ "$_setup_nvidia_usable" = true ]; then - _LINUX_HAS_GPU=true -else - for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo; do - if command -v "$_GPU_TOOL" >/dev/null 2>&1; then - _LINUX_HAS_GPU=true - break - fi - done -fi +# Pick the release repo install_llama_prebuilt.py plans against. Every host this +# installer supports now pulls its llama.cpp prebuilt from the unslothai fork: it +# ships the CUDA (Linux x64/arm64, Windows), ROCm (Linux/Windows) and macOS +# bundles, plus the CPU bundles for Linux/Windows on both x86_64 and arm64. +# ggml-org artifacts are no longer used by default. +_HELPER_RELEASE_REPO="unslothai/llama.cpp" # UNSLOTH_ROCM_GFX_ARCH may be set on a host where no probe fired, so the override # nested in the AMD-detected branch above never ran and _setup_gfx is still empty. -# Honour it here so the routing guard below and the --rocm-gfx forwarding both see -# it (install_llama_prebuilt.py reads the same env var as the --rocm-gfx default). -if [ "$_setup_nvidia_usable" != true ] && [ -z "${_setup_gfx:-}" ] && [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then +# Honour it here so the --rocm-gfx forwarding below still sees it +# (install_llama_prebuilt.py reads the same env var as the --rocm-gfx default). +if [ "${_setup_nvidia_usable:-}" != true ] && [ -z "${_setup_gfx:-}" ] && [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then _setup_gfx="${UNSLOTH_ROCM_GFX_ARCH}" fi -# A resolved/forwarded gfx arch (UNSLOTH_ROCM_GFX_ARCH) means an AMD GPU even when -# no ROCm tooling is on PATH; route it to the fork so the per-gfx prebuilt is -# picked instead of ggml-org / a source build. -if [ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]; then - _LINUX_HAS_GPU=true -fi - -if [ "$_HOST_SYSTEM" = "Linux" ] \ - && [ "$_HOST_MACHINE" = "x86_64" ] \ - && [ "$_LINUX_HAS_GPU" = false ]; then - _HELPER_RELEASE_REPO="ggml-org/llama.cpp" -elif [ "$_HOST_SYSTEM" = "Linux" ] \ - && { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; } \ - && [ "$_LINUX_HAS_GPU" = false ]; then - # CPU-only Linux ARM64 (Ampere Altra, Raspberry Pi 5, GitHub - # `ubuntu-24.04-arm`, CPU-only Jetson rescue mode, ...). The fork ships no - # arm64 CPU bundle, so without this branch the prebuilt resolver returns 0 - # attempts and the installer falls back to a source build. ggml-org ships - # llama-bNNNN-bin-ubuntu-arm64.tar.gz from at least b9072 onward. - _HELPER_RELEASE_REPO="ggml-org/llama.cpp" -else - # GPU Linux (x64 CUDA/ROCm, arm64 CUDA), Windows (CUDA/ROCm), and macOS. - _HELPER_RELEASE_REPO="unslothai/llama.cpp" -fi -unset _GPU_TOOL _LLAMA_PR="${UNSLOTH_LLAMA_PR:-}" _SKIP_PREBUILT_INSTALL=false _LLAMA_PR_FORCE="${UNSLOTH_LLAMA_PR_FORCE:-${_DEFAULT_LLAMA_PR_FORCE}}" @@ -1939,19 +1895,19 @@ else fi # end _SKIP_GGUF_BUILD check # ── arm64 Linux GPU: CPU prebuilt as a last resort ── -# arm64 Linux with a GPU has no CUDA prebuilt anywhere (the unslothai fork is -# x64 only; ggml-org ships no Linux CUDA build), so it source-builds for the -# GPU above. If that produced no binary, install ggml-org's arm64 CPU prebuilt -# instead of leaving the host without llama.cpp. +# An arm64 Linux GPU host source-builds for the GPU above. If that produced no +# binary, install the fork's arm64 CPU prebuilt (app--linux-arm64-cpu.tar.gz) +# instead of leaving the host without llama.cpp. --cpu-fallback drops the GPU +# attributes so the CPU bundle is selected rather than re-attempting CUDA. if [ "$_LLAMA_CPP_DEGRADED" = true ] \ && [ "$_HOST_SYSTEM" = "Linux" ] \ && { [ "$_HOST_MACHINE" = "aarch64" ] || [ "$_HOST_MACHINE" = "arm64" ]; }; then - substep "GPU source build unavailable; trying ggml-org arm64 CPU prebuilt..." + substep "GPU source build unavailable; trying arm64 CPU prebuilt..." _ARM64_CPU_CMD=( python "$SCRIPT_DIR/install_llama_prebuilt.py" --install-dir "$LLAMA_CPP_DIR" --llama-tag "$_REQUESTED_LLAMA_TAG" - --published-repo "ggml-org/llama.cpp" + --published-repo "unslothai/llama.cpp" --cpu-fallback ) # Trust the installer's exit code: it validates the server before exiting 0, diff --git a/tests/studio/install/test_llama_pr_force_and_source.py b/tests/studio/install/test_llama_pr_force_and_source.py index 8f40c9d720..4ff8c349c3 100644 --- a/tests/studio/install/test_llama_pr_force_and_source.py +++ b/tests/studio/install/test_llama_pr_force_and_source.py @@ -362,8 +362,11 @@ class TestSourcePatternsSh: assert '_LLAMA_SOURCE="${_DEFAULT_LLAMA_SOURCE}"' in self.content def test_release_repo_override_removed(self): + # No env-based release-repo override, and CPU-only hosts no longer fall + # back to ggml-org -- every host now routes to the fork. assert "UNSLOTH_LLAMA_RELEASE_REPO:-unslothai/llama.cpp" not in self.content - assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' in self.content + assert '_HELPER_RELEASE_REPO="unslothai/llama.cpp"' in self.content + assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' not in self.content def test_force_compile_skips_prebuilt_resolution_early(self): assert 'if [ "$_LLAMA_FORCE_COMPILE" = "1" ]; then' in self.content @@ -425,12 +428,11 @@ class TestSourcePatternsPs1: assert "$LlamaSource = $DefaultLlamaSource" in self.content def test_release_repo_override_removed(self): - # Repo chosen by GPU detection (GPU -> fork, CPU -> ggml-org), no env override. + # No env-based release-repo override; every host now routes to the fork + # (the CPU-only ggml-org fallback was removed), mirroring setup.sh. assert "$HelperReleaseRepo = if ($env:UNSLOTH_LLAMA_RELEASE_REPO)" not in self.content - assert ( - "$HelperReleaseRepo = if ($HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch) " - '{ "unslothai/llama.cpp" } else { "ggml-org/llama.cpp" }' in self.content - ) + assert '$HelperReleaseRepo = "unslothai/llama.cpp"' in self.content + assert "$HelperReleaseRepo = if (" not in self.content def test_force_compile_skips_prebuilt_resolution_early(self): assert 'if ($env:UNSLOTH_LLAMA_FORCE_COMPILE -eq "1") {' in self.content diff --git a/tests/studio/install/test_pr4562_bugfixes.py b/tests/studio/install/test_pr4562_bugfixes.py index a8cec0cb85..0d2b092924 100644 --- a/tests/studio/install/test_pr4562_bugfixes.py +++ b/tests/studio/install/test_pr4562_bugfixes.py @@ -658,15 +658,21 @@ class TestSourceCodePatterns: assert "_HELPER_RELEASE_REPO}/releases/latest" not in content assert "ggml-org/llama.cpp/releases/latest" not in content - def test_setup_sh_routes_to_fork_only_on_usable_gpu(self): - """Linux routing gates NVIDIA on GPU usability, not nvidia-smi presence, so - CPU-only/hidden-GPU hosts get the ggml CPU prebuilt. Guards the old presence-only loop.""" + def test_setup_sh_routes_every_host_to_fork(self): + """CPU-only Linux (the last ggml-org artifact consumer) now routes to the + fork like every other host, so the release-repo decision is unconditional. + Guards against a silent reintroduction of a ggml-org CPU routing branch. + GPU usability detection (used for PyTorch / source decisions) must stay.""" content = SETUP_SH.read_text() + assert '_HELPER_RELEASE_REPO="unslothai/llama.cpp"' in content + assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' not in content + # Usability gating (not routing) still distinguishes a hidden GPU. assert '[ "$_setup_nvidia_usable" = true ]' in content assert "CUDA_VISIBLE_DEVICES" in content - # nvidia-smi must NOT be back in the bare presence loop. - assert "for _GPU_TOOL in nvidia-smi" not in content - assert "for _GPU_TOOL in rocminfo amd-smi hipconfig hipinfo" in content + # The GPU-tooling probe (PR #4562) stays: ROCm detection goes through + # command -v, not a bare presence loop that mishandled a hidden nvidia-smi. + assert "command -v rocminfo" in content + assert "command -v amd-smi" in content def test_setup_sh_reports_installed_prebuilt_release(self): """Shell wrapper should report the installed prebuilt release from metadata.""" diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index bfc8132683..5cabf41f57 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -3164,21 +3164,25 @@ class TestRocmGfxForwarding: assert "--rocm-gfx" in source assert "$script:ROCmGfxArch" in source - def test_setup_sh_routes_inferred_gfx_to_fork(self): - # An inferred gfx arch must route to the fork even without ROCm tooling. - # Pin the specific guard (a bare "${_setup_gfx:-}" also appears elsewhere). + def test_setup_sh_routes_unconditionally_to_fork(self): + # CPU-only hosts no longer fall back to ggml-org -- the release-repo + # decision is an unconditional fork assignment now. Pin the line text. source = _SETUP_SH_PATH.read_text(encoding = "utf-8") - assert '[ "$_LINUX_HAS_GPU" = false ] && [ -n "${_setup_gfx:-}" ]' in source + assert '_HELPER_RELEASE_REPO="unslothai/llama.cpp"' in source + assert '_HELPER_RELEASE_REPO="ggml-org/llama.cpp"' not in source - def test_setup_ps1_routes_inferred_gfx_to_fork(self): - # On Windows, a resolved $script:ROCmGfxArch counts as a fork install - # even when $HasROCm is false (Adrenalin-only, no HIP runtime). + def test_setup_ps1_routes_unconditionally_to_fork(self): + # Same on Windows: the fork now ships the windows-cpu / windows-arm64 + # bundles, so $HelperReleaseRepo is an unconditional fork assignment. source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") - assert "$HasNvidiaSmi -or $HasROCm -or $script:ROCmGfxArch" in source + assert '$HelperReleaseRepo = "unslothai/llama.cpp"' in source + assert "$HelperReleaseRepo = if (" not in source - # The assertions above pin the guard *text*; the tests below *execute* the - # real routing block and assert the resolved repo, so a refactor that keeps - # the literal but breaks the inferred-gfx -> fork decision is still caught. + # The text pins above guard the literal. The tests below *execute* the real + # routing line from setup.sh / setup.ps1 and assert the resolved release repo, + # so a refactor that reintroduces a conditional (or a ggml-org branch) is still + # caught. Inputs are varied -- CPU-only, inferred/forwarded gfx, usable NVIDIA -- + # to prove no host slips back onto ggml-org. No GPU, no tooling, no network. @staticmethod def _resolve_setup_sh_repo( @@ -3187,15 +3191,18 @@ class TestRocmGfxForwarding: setup_gfx, rocm_gfx_arch_env = "", ): - """Run setup.sh's routing block under bash with PATH emptied (no ROCm tooling) and return _HELPER_RELEASE_REPO.""" + """Run setup.sh's release-repo routing block under bash and return the + resolved _HELPER_RELEASE_REPO. PATH is emptied so any stray tooling probe + misses; routing is unconditional, so the GPU inputs only prove no branch + reroutes a host to ggml-org.""" import shutil bash = shutil.which("bash") if bash is None: pytest.skip("bash not available") source = _SETUP_SH_PATH.read_text(encoding = "utf-8") - start = source.index("\n_LINUX_HAS_GPU=false\n") + 1 - end = source.index("\nunset _GPU_TOOL", start) + len("\nunset _GPU_TOOL") + start = source.index('\n_HELPER_RELEASE_REPO="unslothai/llama.cpp"\n') + 1 + end = source.index("\n_LLAMA_PR=", start) block = source[start:end] assert "_HELPER_RELEASE_REPO" in block, "setup.sh routing anchors not found" env = { @@ -3217,25 +3224,31 @@ class TestRocmGfxForwarding: assert result.returncode == 0, result.stderr return result.stdout.strip() - def test_setup_sh_inferred_gfx_resolves_to_fork(self): - # Only a name-inferred gfx arch -> route to the fork's per-gfx prebuilt - # (not ggml-org). x64 and arm64 share the fork branch. - assert self._resolve_setup_sh_repo("x86_64", False, "gfx1100") == "unslothai/llama.cpp" - assert self._resolve_setup_sh_repo("aarch64", False, "gfx1100") == "unslothai/llama.cpp" - - def test_setup_sh_env_forwarded_gfx_resolves_to_fork(self): - # No probe fired but UNSLOTH_ROCM_GFX_ARCH is set: adopt the env arch - # and route to the fork, same as name-inference. - repo = self._resolve_setup_sh_repo("x86_64", False, "", rocm_gfx_arch_env = "gfx1100") - assert repo == "unslothai/llama.cpp" - - def test_setup_sh_cpu_host_still_resolves_to_ggml(self): - # A real CPU host (no GPU, no inferred gfx, no env override) must keep routing to ggml-org. - assert self._resolve_setup_sh_repo("x86_64", False, "") == "ggml-org/llama.cpp" + @pytest.mark.parametrize( + "machine, nvidia_usable, setup_gfx, env_gfx", + [ + ("x86_64", False, "", ""), # plain CPU host (used to take ggml-org) + ("aarch64", False, "", ""), # plain CPU arm64 host (used to take ggml-org) + ("x86_64", False, "gfx1100", ""), # name-inferred gfx + ("x86_64", False, "", "gfx1100"), # env-forwarded gfx + ("x86_64", True, "", ""), # usable NVIDIA + ], + ) + def test_setup_sh_routing_block_always_resolves_to_fork( + self, machine, nvidia_usable, setup_gfx, env_gfx + ): + assert ( + self._resolve_setup_sh_repo( + machine, nvidia_usable, setup_gfx, rocm_gfx_arch_env = env_gfx + ) + == "unslothai/llama.cpp" + ) @staticmethod - def _resolve_setup_ps1_repo(has_nvidia, has_rocm, gfx_arch): - """Run setup.ps1's $HelperReleaseRepo selection under pwsh and return the resolved repo.""" + def _resolve_setup_ps1_repo(): + """Run setup.ps1's $HelperReleaseRepo assignment under pwsh and return the + resolved repo. The assignment is unconditional now, so there are no host + inputs to vary.""" import shutil pwsh = shutil.which("pwsh") @@ -3243,21 +3256,11 @@ class TestRocmGfxForwarding: pytest.skip("pwsh not available") source = _SETUP_PS1_PATH.read_text(encoding = "utf-8") line = next( - ( - ln - for ln in source.splitlines() - if ln.strip().startswith("$HelperReleaseRepo = if (") - ), + (ln for ln in source.splitlines() if ln.strip().startswith("$HelperReleaseRepo =")), None, ) assert line is not None, "$HelperReleaseRepo selection not found in setup.ps1" - harness = ( - f"$HasNvidiaSmi = ${'true' if has_nvidia else 'false'}\n" - f"$HasROCm = ${'true' if has_rocm else 'false'}\n" - f"$script:ROCmGfxArch = '{gfx_arch}'\n" - f"{line}\n" - "Write-Output $HelperReleaseRepo" - ) + harness = f"{line}\nWrite-Output $HelperReleaseRepo" result = subprocess.run( [pwsh, "-NoProfile", "-Command", harness], capture_output = True, @@ -3267,13 +3270,10 @@ class TestRocmGfxForwarding: assert result.returncode == 0, result.stderr return result.stdout.strip() - def test_setup_ps1_inferred_gfx_resolves_to_fork(self): - # Adrenalin-only host: $HasROCm false but gfx inferred -> fork's windows-rocm bundle. - assert self._resolve_setup_ps1_repo(False, False, "gfx1100") == "unslothai/llama.cpp" - - def test_setup_ps1_cpu_host_still_resolves_to_ggml(self): - # No NVIDIA, no ROCm, no inferred gfx -> CPU host stays on ggml-org. - assert self._resolve_setup_ps1_repo(False, False, "") == "ggml-org/llama.cpp" + def test_setup_ps1_routing_resolves_to_fork(self): + # Windows routing is unconditional now: CPU-only Windows (x64 and arm64) + # uses the fork's windows-cpu / windows-arm64 bundles, not ggml-org. + assert self._resolve_setup_ps1_repo() == "unslothai/llama.cpp" # TEST: _pick_rocm_gfx_target -- visible-device selection from rocminfo output. diff --git a/tests/studio/install/test_selection_logic.py b/tests/studio/install/test_selection_logic.py index 464aa9a578..92f08c646b 100644 --- a/tests/studio/install/test_selection_logic.py +++ b/tests/studio/install/test_selection_logic.py @@ -1720,8 +1720,10 @@ class TestResolveInstallAttempts: assert approved.release_tag == "llama-prebuilt-latest" def test_linux_cpu_fork_without_bundle_raises_no_upstream_fallback(self, monkeypatch): - # CPU-only Linux on the fork must not fall back to the ggml-org CPU asset; with - # no fork CPU bundle the resolver raises rather than reaching upstream. + # A CPU-only Linux host on the fork never falls back to the ggml-org CPU + # asset. CPU-only Linux now routes to the fork, but if a release manifest + # happens to ship no CPU bundle the resolver raises rather than quietly + # reaching for an upstream asset. host = make_host( has_usable_nvidia = False, has_physical_nvidia = False, @@ -1846,6 +1848,151 @@ class TestResolveInstallAttempts: assert attempts[0].name == asset_name assert attempts[0].source_label == "published" + @pytest.mark.parametrize( + "system, machine, asset_name, install_kind, bundle_profile", + [ + # CPU-only Linux x64 -> fork linux-cpu (was ggml-org ubuntu-x64). + ("Linux", "x86_64", "app-b9625-linux-x64-cpu.tar.gz", "linux-cpu", "linux-cpu-x64"), + # CPU-only Linux arm64 -> fork linux-arm64 (was ggml-org ubuntu-arm64). + ( + "Linux", + "aarch64", + "app-b9625-linux-arm64-cpu.tar.gz", + "linux-arm64", + "linux-cpu-arm64", + ), + # CPU-only Windows arm64 -> fork windows-arm64 (was ggml-org win-cpu-arm64). + ( + "Windows", + "arm64", + "app-b9625-windows-arm64-cpu.zip", + "windows-arm64", + "windows-cpu-arm64", + ), + ], + ) + def test_cpu_host_prefers_published_fork_asset( + self, monkeypatch, system, machine, asset_name, install_kind, bundle_profile + ): + # CPU-only hosts now select the fork's CPU bundle from the manifest and + # must never query ggml-org upstream assets. Windows x64 CPU is covered + # separately by test_windows_cpu_prefers_published_asset. + host = make_host( + system = system, + machine = machine, + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + release = make_release( + [ + make_artifact( + asset_name, + install_kind = install_kind, + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = bundle_profile, + rank = 1000, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9625", + assets = {asset_name: f"https://published.example/{asset_name}"}, + ) + checksums = make_checksums_with_source( + [asset_name], + release_tag = release.release_tag, + upstream_tag = "b9625", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "github_release_assets", + lambda repo, tag: (_ for _ in ()).throw( + AssertionError("fork CPU host must not query upstream assets") + ), + ) + + _requested_tag, resolved_tag, attempts, _approved = resolve_install_attempts( + "latest", + host, + "unslothai/llama.cpp", + "", + ) + + assert resolved_tag == "b9625" + assert attempts[0].name == asset_name + assert attempts[0].install_kind == install_kind + assert attempts[0].source_label == "published" + + def test_cpu_only_unsupported_arch_source_builds(self, monkeypatch): + # A CPU-only Linux host that is neither x86_64 nor arm64 (ppc64le, + # riscv64, s390x) has no compatible CPU bundle. It must source-build, not + # receive the x86_64 linux-cpu binary (the Linux preflight checks libs, + # not ELF arch, so a wrong-arch binary would slip through). + host = make_host( + machine = "ppc64le", + has_usable_nvidia = False, + has_physical_nvidia = False, + nvidia_smi = None, + ) + assert not host.is_x86_64 and not host.is_arm64 + x64_asset = "app-b9625-linux-x64-cpu.tar.gz" + release = make_release( + [ + make_artifact( + x64_asset, + install_kind = "linux-cpu", + runtime_line = None, + coverage_class = None, + supported_sms = [], + min_sm = None, + max_sm = None, + bundle_profile = "linux-cpu-x64", + rank = 1000, + ) + ], + release_tag = "llama-prebuilt-latest", + upstream_tag = "b9625", + assets = {x64_asset: f"https://published.example/{x64_asset}"}, + ) + checksums = make_checksums_with_source( + [x64_asset], + release_tag = release.release_tag, + upstream_tag = "b9625", + ) + + monkeypatch.setattr( + INSTALL_LLAMA_PREBUILT, + "iter_resolved_published_releases", + lambda requested_tag, published_repo, published_release_tag = "": iter( + [ + INSTALL_LLAMA_PREBUILT.ResolvedPublishedRelease( + bundle = release, + checksums = checksums, + ) + ] + ), + ) + + with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt asset was found"): + resolve_install_attempts("latest", host, "unslothai/llama.cpp", "") + def test_macos_prefers_published_asset(self, monkeypatch): host = make_host( system = "Darwin", @@ -3447,8 +3594,9 @@ class TestLinuxArm64ForkFallsBackToSource: assert plans == ["plan"] def test_arm64_cpu_on_ggml_org_is_not_blocked(self, monkeypatch): - # CPU-only arm64 routes to ggml-org, so the guard must not fire; it reaches the - # iterator (empty here -> generic message). + # ggml-org is reachable only via an explicit --published-repo override now, + # but the guard must still not fire on arm64 there; it reaches the iterator + # (empty here -> generic message). monkeypatch.setattr( INSTALL_LLAMA_PREBUILT, "iter_release_payloads_by_time", @@ -3474,7 +3622,7 @@ class TestLinuxArm64ForkFallsBackToSource: class TestCpuFallback: - """--cpu-fallback drops GPU attributes so the host's OS/arch CPU prebuilt is selected, letting an arm64 GPU host install ggml-org's arm64 CPU build when its source build produced no binary.""" + """--cpu-fallback drops GPU attributes so the host's OS/arch CPU prebuilt is selected, letting an arm64 GPU host install the fork's arm64 CPU bundle when its source build produced no binary.""" _SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh" @@ -3553,10 +3701,15 @@ class TestCpuFallback: def test_setup_sh_has_arm64_cpu_prebuilt_fallback(self): source = self._SETUP_SH.read_text(encoding = "utf-8") - assert "--cpu-fallback" in source - # Fallback targets ggml-org (only repo with an arm64 Linux build), gated on a - # degraded arm64 source build. - assert "ggml-org/llama.cpp" in source + # The arm64 GPU last-resort CPU fallback now pulls the fork's arm64 CPU + # bundle (app--linux-arm64-cpu.tar.gz), not ggml-org's, and is gated + # on a degraded source build for arm64. + start = source.index("_ARM64_CPU_CMD=(") + end = source.index(")", start) + block = source[start:end] + assert "--cpu-fallback" in block + assert '--published-repo "unslothai/llama.cpp"' in block + assert '--published-repo "ggml-org/llama.cpp"' not in block assert "_LLAMA_CPP_DEGRADED" in source From d0c8d550a6db537583dbb78dd186e56ff103d2fa Mon Sep 17 00:00:00 2001 From: Tai An Date: Wed, 8 Jul 2026 05:38:06 -0700 Subject: [PATCH 025/402] fix(studio/hub): apply repo_id length limit per segment, not whole string (#6946) (#6953) * fix(studio/hub): apply repo_id length limit per segment, not whole string is_valid_repo_id() applied the 96-char limit to the full "namespace/repo_name" string, so a repo with a valid (<=96 char) name but a long combined id was falsely rejected. Match huggingface_hub.validate_repo_id by checking the length per segment instead. Fixes #6946. * Fix long repo id state filenames * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../backend/hub/tests/test_model_services.py | 46 +++++++++++++++++++ studio/backend/hub/utils/paths.py | 9 +++- studio/backend/hub/utils/state_dir.py | 33 +++++++++++-- 3 files changed, 82 insertions(+), 6 deletions(-) diff --git a/studio/backend/hub/tests/test_model_services.py b/studio/backend/hub/tests/test_model_services.py index f05d8359ec..44701c0b64 100644 --- a/studio/backend/hub/tests/test_model_services.py +++ b/studio/backend/hub/tests/test_model_services.py @@ -105,6 +105,10 @@ def test_repo_id_validation_accepts_hf_repo_id_contract(repo_id): assert paths.is_valid_repo_id(repo_id) +def test_repo_id_validation_accepts_max_length_namespaced_repo(): + assert paths.is_valid_repo_id(f"{'a' * 96}/{'b' * 96}") + + @pytest.mark.parametrize( "repo_id", [ @@ -121,6 +125,48 @@ def test_repo_id_validation_rejects_unsafe_or_invalid_ids(repo_id): assert not paths.is_valid_repo_id(repo_id) +def test_download_state_preserves_readable_keys_when_safe(monkeypatch, tmp_path): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + + path = state_dir.marker_path("model", "Owner/Repo", "Q4_K_M") + + assert path is not None + assert path.name == "models--owner--repo--variant--q4_k_m.json" + + +@pytest.mark.parametrize("variant", ["bad variant with spaces", "q" * 64]) +def test_download_state_bounds_long_repo_variant_filenames(monkeypatch, tmp_path, variant): + monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path) + repo_id = f"{'a' * 96}/{'b' * 96}" + + assert paths.is_valid_repo_id(repo_id) + assert download_manifest.write_cancel_marker("model", repo_id, variant, "http") + assert download_manifest.write_manifest( + "model", + repo_id, + variant, + [download_manifest.ExpectedFile(path = "model.gguf", size = 1)], + "http", + ) + + marker_path = state_dir.marker_path("model", repo_id, variant) + manifest_path = state_dir.manifest_path("model", repo_id, variant) + + assert marker_path is not None + assert manifest_path is not None + assert "--sha256-" in marker_path.name + assert len(marker_path.name.encode("utf-8")) <= 255 + assert len(f".{marker_path.name}.tmp-00000000".encode("utf-8")) <= 255 + assert download_manifest.has_cancel_marker("model", repo_id, variant) + assert download_manifest.read_manifest("model", repo_id, variant) is not None + assert list(download_manifest.iter_variant_markers("model", repo_id)) == [ + (variant, marker_path) + ] + assert list(download_manifest.iter_variant_manifests("model", repo_id)) == [ + (variant, manifest_path) + ] + + class _RecordingLogger: def __init__(self): self.warnings = [] diff --git a/studio/backend/hub/utils/paths.py b/studio/backend/hub/utils/paths.py index afcb0b41dc..5435202565 100644 --- a/studio/backend/hub/utils/paths.py +++ b/studio/backend/hub/utils/paths.py @@ -181,15 +181,20 @@ def is_valid_repo_id(repo_id: str) -> bool: """Validate Hugging Face ``repo_name`` or ``namespace/repo_name`` IDs.""" if not repo_id or repo_id != repo_id.strip(): return False - if len(repo_id) > _MAX_REPO_ID_LENGTH or repo_id.endswith(".git"): + if repo_id.endswith(".git"): return False if "--" in repo_id or ".." in repo_id: return False segments = repo_id.split("/") if len(segments) not in (1, 2): return False + # Match huggingface_hub.validate_repo_id: the 96-char limit applies per + # segment (repo name / namespace), not to the whole "namespace/repo_name" + # string, so long-but-valid repo names are not falsely rejected. return all( - segment not in ("", ".", "..") and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None + segment not in ("", ".", "..") + and len(segment) <= _MAX_REPO_ID_LENGTH + and _VALID_REPO_ID_SEGMENT.fullmatch(segment) is not None for segment in segments ) diff --git a/studio/backend/hub/utils/state_dir.py b/studio/backend/hub/utils/state_dir.py index a304477a3d..183e934724 100644 --- a/studio/backend/hub/utils/state_dir.py +++ b/studio/backend/hub/utils/state_dir.py @@ -11,8 +11,9 @@ cache lifecycle. Two subdirectories: manifests/ .json per-download expected-files manifest cancelled/ .json per-download cancel marker -The ```` mirrors HF's cache dir naming so a state file can be -eyeballed next to the on-disk repo it describes: +The ```` mirrors HF's cache dir naming while the resulting manifest, +cancel-marker, and atomic-write temp filenames fit common filesystem basename +limits. Very long repo IDs use a stable hash in the state key: models---- full snapshot models------variant-- GGUF variant @@ -49,6 +50,11 @@ _MANIFESTS_SUBDIR = "manifests" _CANCELLED_SUBDIR = "cancelled" _WORKERS_SUBDIR = "workers" _SAFE_VARIANT_FRAGMENT = re.compile(r"^[a-z0-9._-]{1,64}$") +_MAX_STATE_BASENAME_BYTES = 255 +_STATE_EXTENSION = ".json" +# _atomic_write_json writes "..tmp-<8hex>" beside the final file. +_ATOMIC_WRITE_TMP_OVERHEAD = len(".") + len(".tmp-") + 8 +_MAX_VARIANT_FRAGMENT_LENGTH = 64 def state_root() -> Optional[Path]: @@ -84,16 +90,35 @@ def repo_cache_basename(repo_type: RepoType, repo_id: str) -> str: return f"{repo_type}s--{repo_id.replace('/', '--')}".lower() +def _filename_bytes(name: str) -> int: + return len(name.encode("utf-8")) + + +def _state_filename_fits(entry_key: str) -> bool: + filename = f"{entry_key}{_STATE_EXTENSION}" + return _filename_bytes(filename) + _ATOMIC_WRITE_TMP_OVERHEAD <= _MAX_STATE_BASENAME_BYTES + + +def _state_repo_key(repo_type: RepoType, repo_id: str) -> str: + base = repo_cache_basename(repo_type, repo_id) + variant_prefix = f"{base}--variant--" + longest_variant_key = f"{variant_prefix}{'x' * _MAX_VARIANT_FRAGMENT_LENGTH}" + if _state_filename_fits(longest_variant_key): + return base + digest = hashlib.sha256(base.encode("utf-8")).hexdigest()[:32] + return f"{repo_type}s--sha256-{digest}" + + def variant_filename_prefix(repo_type: RepoType, repo_id: str) -> str: """Lowercased prefix every variant-keyed state file for this repo shares. The single source the download_manifest enumerators match against, so the scheme in :func:`_entry_key` cannot drift from them silently.""" - return f"{repo_cache_basename(repo_type, repo_id)}--variant--" + return f"{_state_repo_key(repo_type, repo_id)}--variant--" def _entry_key(repo_type: RepoType, repo_id: str, variant: Optional[str]) -> str: - base = repo_cache_basename(repo_type, repo_id) + base = _state_repo_key(repo_type, repo_id) if not variant: return base normalized_variant = variant.strip().lower() From 62a6eb2a3df395e2e0e94218cfc963318f5c7392 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 05:57:44 -0700 Subject: [PATCH 026/402] MoE LoRA: auto-target per-expert Linear experts (gpt-oss 4bit) instead of leaving them frozen (#6936) * models: auto-target per-expert Linear MoE experts for LoRA (gpt-oss 4bit) MoE checkpoints whose experts are stored as per-expert nn.Linear ModuleLists could not receive expert LoRA. gpt-oss bnb-4bit is the canonical case: its experts live at mlp.experts.gate_up_projs. and mlp.experts.down_projs. as per-expert Linear4bit modules, not a fused nn.Parameter. The target_parameters path only handles the fused nn.Parameter layout, and the plain gate_proj/up_proj/down_proj leaf names do not match the per-expert indices, so get_peft_model attached LoRA to attention only and left every expert frozen (0 of 1536 on gpt-oss-20b) even though the grouped bnb-4bit training forward exists. Add get_moe_target_modules, the module-LoRA counterpart of get_moe_target_parameters: it detects per-expert Linear ModuleLists under an experts container and returns their suffix target_modules names (gate_up_projs. / down_projs.). get_peft_model in both llama.py and vision.py extends target_modules with these, handling the explicit leaf-list form and the regex form (auto / all-linear / scoped). It is gated on the same MLP-in-scope condition as the parameter path, so an attention-only request still skips the experts. Also gate get_moe_target_parameters on the fused parameter actually existing, so a per-expert-Linear layout no longer produces a dead target_parameters path or a misleading "Enabling LoRA on MoE parameters" line; those experts are handled through target_modules instead. Validated on gpt-oss-20b-unsloth-bnb-4bit (transformers 5.5.0): experts attach (1536 modules, trainable 0.036 percent to 1.65 percent) across the default, None and all-linear paths; training memorizes and the LoRA adapter reproduces exactly after a cold reload in a fresh process. No regression: fused-parameter MoEs (Qwen3-30B-A3B-4bit), non-MoE models, and attention-only requests are unaffected (get_moe_target_modules returns an empty list). Merging these per-expert adapters into a merged_16bit checkpoint is handled by a companion unsloth-zoo change (saving_utils folds each per-expert delta into the fused gate_up_proj / down_proj tensor). With both, the LoRA adapter and the merged_16bit checkpoint reload the trained behavior identically. * models: scope per-expert MoE targets, keep repeat get_peft_model idempotent, warn on old zoo Address review of the per-expert Linear MoE targeting: - Scope get_moe_target_modules to the requested projection leaves (gate/up map to the gate_up ModuleList, down maps to the down ModuleList), so a narrowed request such as target_modules=["down_proj"] no longer also trains gate_up_projs, matching get_moe_target_parameters. - Detect experts through a PEFT-wrapped base_layer as well, and recompute the auto-added expert targets in the llama.py existing-adapter check, so a repeat get_peft_model call with the same arguments stays idempotent instead of raising on the saved expert targets. - Warn when the installed unsloth_zoo cannot fold these per-expert experts into a merged_16bit checkpoint (older releases keep the fused gate_up_proj / down_proj tensors and drop the per-expert deltas), so the expert LoRA is not silently lost on save_pretrained_merged; the fold lands in unsloth-zoo #885. The LoRA adapter itself is unaffected. * [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> --- unsloth/models/_utils.py | 119 +++++++++++++++++++++++++++++++++++++-- unsloth/models/llama.py | 18 ++++++ unsloth/models/vision.py | 29 ++++++++++ 3 files changed, 162 insertions(+), 4 deletions(-) diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index 1c75f8ce66..b68cb702b1 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -86,6 +86,8 @@ __all__ = [ "maybe_prefetch_hf_snapshot", "is_moe_model", "get_moe_target_parameters", + "get_moe_target_modules", + "warn_if_zoo_cannot_merge_moe_experts", "_select_moe_detection_targets", "make_fast_generate_wrapper", "_mark_unsloth_disable_data_parallel", @@ -4060,13 +4062,17 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str alternate_name = "experts.down_proj", ) - # gate_up_proj combines both gate_proj and up_proj in MoE - # Also match "gate_up_proj" directly since users may specify the fused name + # gate_up_proj combines gate_proj and up_proj; also match the fused name directly. + # Only target a fused expert Parameter that exists: per-expert Linear layouts + # (e.g. gpt-oss bnb-4bit) have no fused Parameter and are handled by + # get_moe_target_modules, so skip them rather than pass PEFT a dead path. if "gate_proj" in target_set or "up_proj" in target_set or "gate_up_proj" in target_set: - moe_params.append(gate_up_name) + if _moe_parameter_exists(model, gate_up_name): + moe_params.append(gate_up_name) if "down_proj" in target_set: - moe_params.append(down_name) + if _moe_parameter_exists(model, down_name): + moe_params.append(down_name) if moe_params: print( @@ -4077,6 +4083,111 @@ def get_moe_target_parameters(model, target_modules = None) -> Optional[List[str return None +def _moe_parameter_exists(model, name: str) -> bool: + """True if ``name`` is an exact suffix of some parameter path on the model.""" + if not hasattr(model, "named_parameters"): + return False + try: + for parameter_name, _ in model.named_parameters(): + if parameter_name == name or parameter_name.endswith("." + name): + return True + except Exception: + return False + return False + + +def get_moe_target_modules(model, target_modules = None) -> List[str]: + """Per-expert ``target_modules`` suffixes for MoE models whose experts are stored + as per-expert ``nn.Linear`` ModuleLists rather than fused nn.Parameters. + + gpt-oss bnb-4bit is the canonical case (mlp.experts.gate_up_projs. / + down_projs. as Linear4bit): no fused Parameter, and the plain + gate/up/down_proj leaves do not match, so LoRA skips them. Returning the + per-expert suffixes makes PEFT attach via ordinary suffix matching (the + module-LoRA counterpart of get_moe_target_parameters). Returns [] for non-MoE, + fused-parameter MoEs, an absent per-expert layout, or a request that omits the + MLP experts (so an attention-only run does not train experts). + """ + if not is_moe_model(model): + return [] + if target_modules is None: + return [] + if isinstance(target_modules, str): + target_set = _moe_target_set_from_string(target_modules) + else: + target_set = { + target + for target in target_modules or () + if (isinstance(target, str) and "." not in target and target in _MOE_BROAD_MLP_TARGETS) + } + if not (target_set & _MOE_BROAD_MLP_TARGETS): + return [] + + if not hasattr(model, "named_modules"): + return [] + + # Scope the returned suffixes to the requested projection leaves, matching + # get_moe_target_parameters: gate_proj/up_proj/gate_up_proj map to the fused + # gate_up ModuleList (e.g. gate_up_projs); down_proj maps to the down ModuleList + # (e.g. down_projs). A down-only (or gate/up-only) request must not pull in the + # other projection. + want_gate_up = bool(target_set & {"gate_proj", "up_proj", "gate_up_proj"}) + want_down = "down_proj" in target_set + + targets = set() + for name, module in model.named_modules(): + if not isinstance(module, torch.nn.ModuleList) or len(module) == 0: + continue + parent, _, leaf = name.rpartition(".") + # ModuleList directly under an ``experts`` container, holding only Linear + # leaves (bnb Linear4bit / Linear8bitLt subclass nn.Linear). After PEFT has + # wrapped the experts the child is a LoRA layer whose ``base_layer`` is the + # Linear, so accept that too (keeps this idempotent across a re-wrapped model). + if not parent.endswith("experts"): + continue + if not all( + isinstance(child, torch.nn.Linear) + or isinstance(getattr(child, "base_layer", None), torch.nn.Linear) + for child in module + ): + continue + # Honor the requested subset: classify the ModuleList by projection role. + leaf_lower = leaf.lower() + is_down = "down" in leaf_lower + is_gate_up = (not is_down) and ("gate" in leaf_lower or "up" in leaf_lower) + if is_down and not want_down: + continue + if is_gate_up and not want_gate_up: + continue + # One entry per expert index; ``leaf.`` matches expert i in every layer. + for expert_index in range(len(module)): + targets.add(f"{leaf}.{expert_index}") + + return sorted(targets) + + +def warn_if_zoo_cannot_merge_moe_experts(): + """Warn once when the installed unsloth_zoo cannot fold per-expert Linear MoE LoRA + into a merged_16bit checkpoint. Older zoo releases keep the fused gate_up_proj / + down_proj tensors and drop the per-expert gate_up_projs. / down_projs. deltas, + so save_pretrained_merged("merged_16bit") would silently lose the expert training + (the LoRA adapter itself still saves and reloads correctly).""" + try: + from unsloth_zoo import saving_utils as _saving_utils + + # _fold_perexpert_lora_into_fused is the helper that folds these experts. + if hasattr(_saving_utils, "_fold_perexpert_lora_into_fused"): + return + except Exception: + return # cannot introspect zoo -> stay quiet rather than false-alarm + logger.warning_once( + "Unsloth: the installed unsloth_zoo will not fold these per-expert experts into " + "a merged_16bit checkpoint, so save_pretrained_merged('merged_16bit') would drop " + "the expert LoRA. Upgrade unsloth_zoo to merge them; saving the LoRA adapter is " + "unaffected." + ) + + def _select_moe_detection_targets( original_target_modules, scoped_target_modules, diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index a1da099758..1f43f61443 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -3105,6 +3105,11 @@ class FastLlamaModel: new_target_modules = list(target_modules) + list( modules_to_save if modules_to_save is not None else [] ) + # Per-expert Linear MoE experts (e.g. gpt-oss bnb-4bit) were auto-added to the + # saved target_modules when the adapter was first created. Recompute them so a + # repeat get_peft_model call with the same args stays idempotent instead of + # tripping the mismatch below. No-op for non per-expert-Linear models. + new_target_modules += get_moe_target_modules(model, target_modules) # Now check! new_target_modules = set(new_target_modules) @@ -3331,6 +3336,19 @@ class FastLlamaModel: if target_parameters is None: target_parameters = get_moe_target_parameters(model, target_modules) + # Per-expert Linear expert layouts (e.g. gpt-oss bnb-4bit) are Linear modules, + # not fused Parameters, so target them via target_modules. No-op otherwise. + _moe_module_targets = get_moe_target_modules(model, target_modules) + if _moe_module_targets: + _added = [t for t in _moe_module_targets if t not in final_modules] + final_modules.extend(_added) + if _added: + print( + f"Unsloth: Detected MoE model with per-expert Linear experts. " + f"Enabling LoRA on {len(_added)} expert projection modules." + ) + warn_if_zoo_cannot_merge_moe_experts() + if finetune_last_n_layers is not None and layers_to_transform is None: from .vision import _get_total_transformer_layers _total_layers = _get_total_transformer_layers(model) diff --git a/unsloth/models/vision.py b/unsloth/models/vision.py index 0a68a49fee..0235d80e93 100644 --- a/unsloth/models/vision.py +++ b/unsloth/models/vision.py @@ -1807,6 +1807,35 @@ class FastBaseModel: ) target_parameters = get_moe_target_parameters(model, _moe_targets) + # Per-expert Linear expert layouts (e.g. gpt-oss bnb-4bit) target experts via + # target_modules, not fused Parameters. Extend either form PEFT accepts: a leaf + # list (explicit) or a regex string (auto / all-linear / scoped). No-op otherwise. + _moe_module_detect = _select_moe_detection_targets( + _moe_detect_target, + target_modules, + finetune_mlp_modules = finetune_mlp_modules, + finetune_language_layers = finetune_language_layers, + ) + _moe_module_targets = get_moe_target_modules(model, _moe_module_detect) + if _moe_module_targets: + if isinstance(target_modules, (list, tuple)): + target_modules = list(target_modules) + [ + target for target in _moe_module_targets if target not in target_modules + ] + elif isinstance(target_modules, str): + _expert_leaves = sorted({t.rsplit(".", 1)[0] for t in _moe_module_targets}) + _expert_alt = ( + r".*\.experts\.(?:" + + "|".join(re.escape(leaf) for leaf in _expert_leaves) + + r")\.\d+" + ) + target_modules = f"(?:{target_modules})|(?:{_expert_alt})" + print( + f"Unsloth: Detected MoE model with per-expert Linear experts. " + f"Enabling LoRA on {len(_moe_module_targets)} expert projection modules." + ) + warn_if_zoo_cannot_merge_moe_experts() + if finetune_last_n_layers is not None and layers_to_transform is None: _total_layers = _get_total_transformer_layers(model) if _total_layers is not None and _total_layers > 0: From 03cbe211a38b78448e5844f25f49b65a449d8b10 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Eric=20=20=F0=9F=87=A7=F0=9F=87=B7?= Date: Wed, 8 Jul 2026 10:38:10 -0300 Subject: [PATCH 027/402] Studio: fix flash-attn and torchao install on Blackwell (sm_100+) GPUs (Closes #6961) (#6970) * fix: Remove moot has_blackwell_gpu() function Fixes unslothai/unsloth#6961. This function skipped flash-attn on Blackwell GPUs because no prebuilt wheel existed; Dao-AILab now ships one and url_exists() already gates resolution. Co-Authored-By: Claude Opus 4.8 * fix: use torchao 0.17.0 for Blackwell Fixes #6961. Torchao 0.16.0's cpp extensions are built against CUDA 12, so on a CUDA-13 torch (cu130 / Blackwell) they fail to load with "libcudart.so.12: cannot open shared object file". Select 0.17.0 there instead: its cpp targets torch 2.11, so it is skipped cleanly rather than crashing. CUDA-12 / ROCm / CPU torch 2.10 keeps 0.16.0 and its working kernels. Co-Authored-By: Claude Opus 4.8 * Condense torchao version-selection comments (no behavior change) * Support torch 2.11 in the Studio installer via the torch2.10 prebuilt wheels Map torch 2.11 to the torch2.10 prebuilt wheels for flash-attn, causal-conv1d, and mamba through wheel_utils.prebuilt_wheel_torch_mm, applied in direct_wheel_url (filename) and flash_attn_wheel_url (version). Those torch2.10 CUDA wheels load and pass each project's own test suite on torch 2.11 (verified on B200), so a torch 2.11 environment gets the prebuilt accelerators instead of skipping or building from source. Raise _CUDA_TORCH_PKG_SPEC to <2.12.0 (torchvision <0.27.0, torchaudio <2.12.0) so the CUDA torch repair path can install torch 2.11, where torchao 0.17's cpp kernels load cleanly. Add tests for the mapping. * Keep has_blackwell_gpu as a False stub for future arch gating * Restore has_blackwell_gpu as a return-False probe kept for future arch gating Keep the nvidia-smi compute_cap detection and its two call sites, but short-circuit with return False at the top so flash-attn is no longer skipped on Blackwell (sm_100+ now has prebuilt wheels and url_exists gates resolution). Drop the early return to re-enable arch-based detection later. --------- Co-authored-by: Claude Opus 4.8 Co-authored-by: Daniel Han --- .../tests/test_mlx_training_worker_config.py | 1 - studio/backend/tests/test_torchao_select.py | 17 +- .../tests/test_training_worker_flash_attn.py | 23 -- studio/backend/utils/wheel_utils.py | 26 ++- studio/install_python_stack.py | 70 +++--- .../test_flash_attn_install_python_stack.py | 203 ++++-------------- 6 files changed, 123 insertions(+), 217 deletions(-) diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 14fc0933d0..503bae8da3 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -37,7 +37,6 @@ def _load_worker_module(): for name in ( "direct_wheel_url", "flash_attn_wheel_url", - "has_blackwell_gpu", "install_wheel", "probe_torch_wheel_env", "url_exists", diff --git a/studio/backend/tests/test_torchao_select.py b/studio/backend/tests/test_torchao_select.py index 2d3dc5fbff..e4775a10a6 100644 --- a/studio/backend/tests/test_torchao_select.py +++ b/studio/backend/tests/test_torchao_select.py @@ -32,16 +32,23 @@ def _load_module(monkeypatch): @pytest.mark.parametrize( "torch_version, expected", [ - # torch 2.10 (the reported bug: cu130 resolves 2.10.0) -> 0.16.0, - # independent of the local +cuXXX/+rocm/+cpu suffix or patch level. - ("2.10.0+cu130", "torchao==0.16.0"), + # torch 2.10 on CUDA <= 12 -> 0.16.0 (its cpp is built for torch 2.10.0 and + # loads against the CUDA-12 PyPI wheel). Independent of patch level. + ("2.10.0+cu128", "torchao==0.16.0"), + ("2.10.0+cu126", "torchao==0.16.0"), ("2.10.0+rocm6.4", "torchao==0.16.0"), ("2.10.0+cpu", "torchao==0.16.0"), ("2.10.1", "torchao==0.16.0"), ("2.10.0", "torchao==0.16.0"), - # Pre-release / dev / rc builds: the minor is cleaned of non-digits. + # torch 2.10 on CUDA >= 13 (Blackwell / cu130): 0.16.0's CUDA-12 cpp can't + # load against a CUDA-13 torch (libcudart.so.12 error), so use 0.17.0. + ("2.10.0+cu130", "torchao==0.17.0"), + ("2.10.0+cu140", "torchao==0.17.0"), + # Pre-release / dev / rc builds: the minor is cleaned of non-digits; the + # CUDA tag still decides 0.16.0 vs 0.17.0. ("2.10.0rc1", "torchao==0.16.0"), - ("2.10.0.dev20250804+cu130", "torchao==0.16.0"), + ("2.10.0.dev20250804+cu130", "torchao==0.17.0"), + ("2.10.0.dev20250804+cu128", "torchao==0.16.0"), ("2.10rc1", "torchao==0.16.0"), # torch 2.11 (reachable via ROCm rocm7.2) and forward -> 0.17.0. ("2.11.0+cu130", "torchao==0.17.0"), diff --git a/studio/backend/tests/test_training_worker_flash_attn.py b/studio/backend/tests/test_training_worker_flash_attn.py index 3c5d6cd094..7e7fc1af48 100644 --- a/studio/backend/tests/test_training_worker_flash_attn.py +++ b/studio/backend/tests/test_training_worker_flash_attn.py @@ -59,7 +59,6 @@ def test_runtime_flash_attn_prefers_prebuilt_wheel(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -88,7 +87,6 @@ def test_runtime_flash_attn_falls_back_to_pypi(monkeypatch): statuses: list[str] = [] monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: False) monkeypatch.setattr(builtins, "__import__", _missing_flash_attn_import()) monkeypatch.setattr( worker, @@ -141,27 +139,6 @@ def test_runtime_flash_attn_skip_env_avoids_all_install_work(monkeypatch): worker._sp.run.assert_not_called() -def test_runtime_flash_attn_skips_on_blackwell(monkeypatch): - statuses: list[str] = [] - install_mock = mock.Mock() - - monkeypatch.delenv(worker._FLASH_ATTN_SKIP_ENV, raising = False) - monkeypatch.setattr(worker, "_should_try_runtime_flash_attn_install", lambda max_seq: True) - monkeypatch.setattr(worker, "has_blackwell_gpu", lambda: True) - monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) - monkeypatch.setattr( - worker, - "_send_status", - lambda queue, message: statuses.append(message), - ) - - worker._ensure_flash_attn_for_long_context(event_queue = [], max_seq_length = 65536) - - install_mock.assert_not_called() - assert len(statuses) == 1 - assert "Blackwell" in statuses[0] - - def test_causal_conv1d_fast_path_preserves_wheel_first_install_args(monkeypatch): install_mock = mock.Mock(return_value = True) monkeypatch.setattr(worker, "_install_package_wheel_first", install_mock) diff --git a/studio/backend/utils/wheel_utils.py b/studio/backend/utils/wheel_utils.py index 98697df83c..1b5926fd49 100644 --- a/studio/backend/utils/wheel_utils.py +++ b/studio/backend/utils/wheel_utils.py @@ -26,11 +26,14 @@ FLASH_ATTN_RELEASE_BASE_URL = "https://github.com/Dao-AILab/flash-attention/rele def has_blackwell_gpu() -> bool: """Return True if any visible NVIDIA GPU has compute capability >= 10.0 (Blackwell). - Dao-AILab ships no flash-attention wheels for these archs and older-arch wheels - fail to load, so callers use this to skip the flash-attn install path. Cached - for the process lifetime; tests mocking nvidia-smi must call + Cached for the process lifetime; tests mocking nvidia-smi must call ``has_blackwell_gpu.cache_clear()`` first. """ + # Detection disabled for now: Dao-AILab ships Blackwell (sm_100+) flash-attn + # wheels and url_exists() already gates resolution, so we no longer skip + # flash-attn on Blackwell. The nvidia-smi probe below is kept for possible + # future arch-based gating; drop this early return to re-enable it. + return False exe = shutil.which("nvidia-smi") if not exe: return False @@ -117,6 +120,19 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non return env +# torch 2.11 has no native prebuilt wheels for flash-attn / causal-conv1d / mamba +# yet, but their torch 2.10 CUDA wheels load and pass the projects' own test suites +# on torch 2.11 (verified on B200: FA2 fwd/bwd, causal-conv1d, and mamba selective +# scan all match reference). Reuse the torch 2.10 wheels on torch 2.11 so a 2.11 +# install still gets these prebuilt accelerators instead of building from source. +_PREBUILT_WHEEL_TORCH_MM = {"2.11": "2.10"} + + +def prebuilt_wheel_torch_mm(torch_mm: str) -> str: + """Map a torch major.minor to the one whose prebuilt accelerator wheels to use.""" + return _PREBUILT_WHEEL_TORCH_MM.get(torch_mm, torch_mm) + + def direct_wheel_url( *, filename_prefix: str, @@ -130,7 +146,7 @@ def direct_wheel_url( filename = ( f"{filename_prefix}-{package_version}" - f"+cu{env['cuda_major']}torch{env['torch_mm']}" + f"+cu{env['cuda_major']}torch{prebuilt_wheel_torch_mm(env['torch_mm'])}" f"cxx11abi{env['cxx11abi']}-{env['python_tag']}-{env['python_tag']}" f"-{env['platform_tag']}.whl" ) @@ -152,7 +168,7 @@ def flash_attn_package_version(torch_mm: str) -> str | None: def flash_attn_wheel_url(env: dict[str, str] | None) -> str | None: if env is None: return None - package_version = flash_attn_package_version(env["torch_mm"]) + package_version = flash_attn_package_version(prebuilt_wheel_torch_mm(env["torch_mm"])) if package_version is None: return None return direct_wheel_url( diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index e033a56a0a..19c492deaa 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -103,33 +103,47 @@ _PYTORCH_WHL_BASE = ( os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" ).rstrip("/") -# CUDA torch repair specs (see _ensure_cuda_torch). torchvision/torchaudio are -# pinned to the torch<2.11 family rather than left bare: the install uses an -# exclusive --index-url (no PyPI fallback), so a bare name could resolve a -# torchvision built against a different torch major (e.g. 0.27 for torch 2.12) -# and fail at runtime with an ABI mismatch. Same bounds as the _default ROCm -# spec above, which targets the same torch family. +# CUDA torch repair specs (see _ensure_cuda_torch). torch 2.11 is allowed: its +# torchao 0.17 cpp kernels load cleanly (0.16 crashes on cu130), and the flash-attn +# / causal-conv1d / mamba torch2.10 wheels load and pass their upstream suites on +# 2.11 (see wheel_utils._PREBUILT_WHEEL_TORCH_MM). torchvision/torchaudio are pinned +# (not bare) because the install uses an exclusive --index-url (no PyPI fallback), so +# a bare name could resolve one built against a different torch major (e.g. 0.27 for +# torch 2.12) and fail at runtime with an ABI mismatch. _CUDA_TORCH_PKG_SPEC: tuple[str, str, str] = ( - "torch>=2.4,<2.11.0", - "torchvision>=0.19,<0.26.0", - "torchaudio>=2.4,<2.11.0", + "torch>=2.4,<2.12.0", + "torchvision>=0.19,<0.27.0", + "torchaudio>=2.4,<2.12.0", ) -# torchao's C++ extensions are built against ONE exact torch release; a newer -# torch makes torchao skip its cpp kernels ("Skipping import of cpp extensions -# due to incompatible torch version ...") and fall back to slow Python. Because -# the torch pin above is a range (and every CUDA index now tops out at torch -# 2.10), the torch actually installed drifts ahead of a fixed torchao pin. So -# pick the torchao whose build matches the torch in the venv. Table: pytorch/ao#2919. -# torch 2.9.x -> torchao 0.14.0 (today's pin; built for torch 2.9.0) -# torch 2.10.x -> torchao 0.16.0 (built for torch 2.10.0) -# torch 2.11.x -> torchao 0.17.0 (built for torch 2.11.0; reachable via ROCm rocm7.2) -# Unknown/older torch keeps the conservative default (no regression vs today). +# torchao's cpp extensions are pinned to ONE torch release AND CUDA major. A torch +# mismatch just skips the cpp kernels (slow Python fallback); a CUDA mismatch fails +# to import ("libcudart.so.12: cannot open shared object file"). The torch pin is a +# range, so match torchao to the installed torch (table: pytorch/ao#2919): +# 2.9.x -> 0.14.0 +# 2.10.x, CUDA<=12 -> 0.16.0 (cpp built for 2.10, loads via the CUDA-12 wheel) +# 2.10.x, CUDA>=13 -> 0.17.0 (cu130: 0.16.0's CUDA-12 cpp crashes on load; 0.17.0 +# targets torch 2.11 so its cpp is cleanly skipped, not crashed) +# 2.11.x -> 0.17.0 (reachable via CUDA or ROCm rocm7.2) +# Unknown/older torch keeps the conservative default. _TORCHAO_DEFAULT_SPEC = "torchao==0.14.0" -_TORCHAO_BY_TORCH_MINOR: dict[int, str] = { - 10: "torchao==0.16.0", - 11: "torchao==0.17.0", -} +_TORCHAO_TORCH_210_SPEC = "torchao==0.16.0" +_TORCHAO_TORCH_210_CUDA13_SPEC = "torchao==0.17.0" +_TORCHAO_TORCH_211_PLUS_SPEC = "torchao==0.17.0" +# torch 2.10 built against CUDA >= this major can't load 0.16.0's CUDA-12 cpp. +_TORCHAO_CUDA13_MIN_MAJOR = 13 + + +def _cuda_major_from_torch_version(torch_version: str) -> int | None: + """Extract the CUDA major from a torch local version tag, e.g. '2.10.0+cu130' + -> 13, '2.10.0+cu128' -> 12. Returns None for rocm/cpu/tagless builds.""" + local = str(torch_version).split("+", 1) + if len(local) < 2 or not local[1].startswith("cu"): + return None + digits = re.sub(r"[^0-9].*", "", local[1][2:]) # 'cu130' -> '130' + if not digits: + return None + return int(digits) // 10 # '130' -> 13, '128' -> 12, '118' -> 11 def _select_torchao_spec(torch_version: str | None) -> str: @@ -151,8 +165,14 @@ def _select_torchao_spec(torch_version: str | None) -> str: if major != 2: return _TORCHAO_DEFAULT_SPEC if minor >= 11: - return _TORCHAO_BY_TORCH_MINOR[11] # newest known build; covers 2.11+ - return _TORCHAO_BY_TORCH_MINOR.get(minor, _TORCHAO_DEFAULT_SPEC) + return _TORCHAO_TORCH_211_PLUS_SPEC # newest known build; covers 2.11+ + if minor == 10: + # cu130+ can't load 0.16.0's CUDA-12 cpp; use 0.17.0 (cpp skipped, not crashed). + cuda_major = _cuda_major_from_torch_version(str(torch_version)) + if cuda_major is not None and cuda_major >= _TORCHAO_CUDA13_MIN_MAJOR: + return _TORCHAO_TORCH_210_CUDA13_SPEC + return _TORCHAO_TORCH_210_SPEC + return _TORCHAO_DEFAULT_SPEC def _probe_installed_torch_version() -> str | None: diff --git a/tests/python/test_flash_attn_install_python_stack.py b/tests/python/test_flash_attn_install_python_stack.py index 26ff03505a..bf3ed57788 100644 --- a/tests/python/test_flash_attn_install_python_stack.py +++ b/tests/python/test_flash_attn_install_python_stack.py @@ -13,102 +13,35 @@ sys.path.insert(0, str(STUDIO_DIR)) sys.path.insert(0, str(STUDIO_DIR / "backend")) import install_python_stack as ips -from backend.utils import wheel_utils +from utils import wheel_utils -def _smi_result(stdout: str, returncode: int = 0) -> subprocess.CompletedProcess: - return subprocess.CompletedProcess(["nvidia-smi"], returncode, stdout, "") +class TestPrebuiltWheelTorchMapping: + def test_torch_211_maps_to_torch210(self): + assert wheel_utils.prebuilt_wheel_torch_mm("2.11") == "2.10" + def test_other_versions_pass_through(self): + for torch_mm in ("2.9", "2.10", "2.12"): + assert wheel_utils.prebuilt_wheel_torch_mm(torch_mm) == torch_mm -class TestHasBlackwellGpu: - def setup_method(self): - wheel_utils.has_blackwell_gpu.cache_clear() - - def teardown_method(self): - wheel_utils.has_blackwell_gpu.cache_clear() - - def test_returns_false_when_nvidia_smi_missing(self): - with mock.patch.object(wheel_utils.shutil, "which", return_value = None): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_true_for_sm_100(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("10.0\n")), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_true_for_sm_120(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.0\n")), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_true_for_sm_121(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("12.1\n")), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_false_for_sm_90(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("9.0\n")), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_false_for_sm_89(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object(wheel_utils.subprocess, "run", return_value = _smi_result("8.9\n")), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_mixed_gpus_with_one_blackwell_returns_true(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - return_value = _smi_result("8.0\n10.0\n"), - ), - ): - assert wheel_utils.has_blackwell_gpu() is True - - def test_returns_false_when_nvidia_smi_fails(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - return_value = _smi_result("", returncode = 1), - ), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_false_on_subprocess_timeout(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - side_effect = subprocess.TimeoutExpired(cmd = "nvidia-smi", timeout = 10), - ), - ): - assert wheel_utils.has_blackwell_gpu() is False - - def test_returns_false_on_malformed_output(self): - with ( - mock.patch.object(wheel_utils.shutil, "which", return_value = "/usr/bin/nvidia-smi"), - mock.patch.object( - wheel_utils.subprocess, - "run", - return_value = _smi_result("not-a-number\n\n"), - ), - ): - assert wheel_utils.has_blackwell_gpu() is False + def test_direct_wheel_url_reuses_torch210_on_211(self): + # causal-conv1d / mamba go through direct_wheel_url; torch 2.11 reuses the + # torch2.10 wheel filename just like flash-attn does. + url = wheel_utils.direct_wheel_url( + filename_prefix = "causal_conv1d", + package_version = "1.6.1", + release_tag = "v1.6.1.post4", + release_base_url = "https://example.test/download", + env = { + "python_tag": "cp313", + "torch_mm": "2.11", + "cuda_major": "13", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + }, + ) + assert url is not None + assert "causal_conv1d-1.6.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url class TestFlashAttnWheelSelection: @@ -118,9 +51,24 @@ class TestFlashAttnWheelSelection: def test_torch_29_maps_to_v283(self): assert ips._select_flash_attn_version("2.9") == "2.8.3" - def test_unsupported_torch_has_no_wheel_mapping(self): + def test_torch_211_has_no_native_version_entry(self): + # The raw version table has no torch2.11-tagged wheel; the URL builder + # reuses the torch2.10 wheel instead (see test_torch_211_reuses_torch210_wheel). assert ips._select_flash_attn_version("2.11") is None + def test_torch_211_reuses_torch210_wheel(self): + url = ips._build_flash_attn_wheel_url( + { + "python_tag": "cp313", + "torch_mm": "2.11", + "cuda_major": "13", + "cxx11abi": "TRUE", + "platform_tag": "linux_x86_64", + } + ) + assert url is not None + assert "flash_attn-2.8.1+cu13torch2.10cxx11abiTRUE-cp313-cp313-linux_x86_64.whl" in url + def test_exact_wheel_url_uses_full_env_tuple(self): url = ips._build_flash_attn_wheel_url( { @@ -333,83 +281,22 @@ class TestEnsureFlashAttn: mock_probe.assert_not_called() mock_install_wheel.assert_not_called() - def test_blackwell_gpu_skips_install_with_warning(self): - step_messages: list[tuple[str, str]] = [] - - def fake_step( - label: str, - value: str, - color_fn = None, - ): - step_messages.append((label, value)) - - with ( - mock.patch.object(ips, "NO_TORCH", False), - mock.patch.object(ips, "IS_WINDOWS", False), - mock.patch.object(ips, "IS_MACOS", False), - mock.patch.object(ips, "has_blackwell_gpu", return_value = True), - mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, - mock.patch.object(ips, "install_wheel") as mock_install_wheel, - mock.patch.object(ips, "_step", side_effect = fake_step), - mock.patch("subprocess.run", return_value = self._import_check()), - ): - ips._ensure_flash_attn() - - mock_probe.assert_not_called() - mock_install_wheel.assert_not_called() - assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages) - - def test_blackwell_gpu_on_windows_emits_blackwell_warning(self): - step_messages: list[tuple[str, str]] = [] - - def fake_step( - label: str, - value: str, - color_fn = None, - ): - step_messages.append((label, value)) - + def test_windows_skips_install_without_probing(self): + # flash-attn is Linux-only: on Windows the installer returns before + # probing the torch env or resolving a wheel (no Windows wheels are + # published upstream). with ( mock.patch.object(ips, "NO_TORCH", False), mock.patch.object(ips, "IS_WINDOWS", True), mock.patch.object(ips, "IS_MACOS", False), - mock.patch.object(ips, "has_blackwell_gpu", return_value = True), mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, mock.patch.object(ips, "install_wheel") as mock_install_wheel, - mock.patch.object(ips, "_step", side_effect = fake_step), mock.patch("subprocess.run", return_value = self._import_check()), ): ips._ensure_flash_attn() mock_probe.assert_not_called() mock_install_wheel.assert_not_called() - assert any(label == "warning" and "Blackwell" in msg for label, msg in step_messages) - - def test_non_blackwell_windows_does_not_emit_blackwell_warning(self): - step_messages: list[tuple[str, str]] = [] - - def fake_step( - label: str, - value: str, - color_fn = None, - ): - step_messages.append((label, value)) - - with ( - mock.patch.object(ips, "NO_TORCH", False), - mock.patch.object(ips, "IS_WINDOWS", True), - mock.patch.object(ips, "IS_MACOS", False), - mock.patch.object(ips, "has_blackwell_gpu", return_value = False), - mock.patch.object(ips, "probe_torch_wheel_env") as mock_probe, - mock.patch.object(ips, "install_wheel") as mock_install_wheel, - mock.patch.object(ips, "_step", side_effect = fake_step), - mock.patch("subprocess.run", return_value = self._import_check()), - ): - ips._ensure_flash_attn() - - mock_probe.assert_not_called() - mock_install_wheel.assert_not_called() - assert not any("Blackwell" in msg for _, msg in step_messages) class TestInstallPythonStackFlashAttnIntegration: From 38ea267124cc3c5a82b96fe9a7a140e4f4e566cc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 06:51:58 -0700 Subject: [PATCH 028/402] Versioning --- pyproject.toml | 6 +++--- unsloth/models/_utils.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 80b3d757e3..2b79121c82 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -73,7 +73,7 @@ triton = [ ] huggingfacenotorch = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "wheel>=0.42.0", "packaging", "numpy", @@ -94,7 +94,7 @@ huggingfacenotorch = [ ] huggingface = [ "unsloth[huggingfacenotorch]", - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "torchvision", "unsloth[triton]", ] @@ -579,7 +579,7 @@ colab-ampere-torch220 = [ "flash-attn>=2.6.3 ; ('linux' in sys_platform)", ] colab-new = [ - "unsloth_zoo>=2026.7.1", + "unsloth_zoo>=2026.7.2", "packaging", "tyro", "transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0", diff --git a/unsloth/models/_utils.py b/unsloth/models/_utils.py index b68cb702b1..fa0e0b1c49 100644 --- a/unsloth/models/_utils.py +++ b/unsloth/models/_utils.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -__version__ = "2026.7.1" +__version__ = "2026.7.2" __all__ = [ "SUPPORTS_BFLOAT16", From 3d41e5868d8aa46ece3b58de63f89be0a5d1d6b9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:22:54 -0700 Subject: [PATCH 029/402] Add has_blackwell_gpu to the mlx worker test's wheel_utils stub (#6980) worker.py imports has_blackwell_gpu from utils.wheel_utils, but _load_worker_module stubs utils.wheel_utils with a fixed name tuple that omitted it, so loading the worker raised ImportError (cannot import name 'has_blackwell_gpu') and Backend CI could not collect test_mlx_training_worker_config.py. Add the name to the stub so it matches worker.py's imports. --- studio/backend/tests/test_mlx_training_worker_config.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/tests/test_mlx_training_worker_config.py b/studio/backend/tests/test_mlx_training_worker_config.py index 503bae8da3..14fc0933d0 100644 --- a/studio/backend/tests/test_mlx_training_worker_config.py +++ b/studio/backend/tests/test_mlx_training_worker_config.py @@ -37,6 +37,7 @@ def _load_worker_module(): for name in ( "direct_wheel_url", "flash_attn_wheel_url", + "has_blackwell_gpu", "install_wheel", "probe_torch_wheel_env", "url_exists", From 116ce48c1a9d2b17596d66a247fe441551bdca91 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:26:10 -0700 Subject: [PATCH 030/402] Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device (#6979) * Studio: allow CPU-only DiffusionGemma by granting the diffusion runner the CPU device * Studio: mark CPU-only DiffusionGemma as non-GPU-resident for training VRAM preflight * Studio: keep the CPU DiffusionGemma change minimal (revert VRAM-flag tweak; Metal hosts still hold unified memory) * Studio: keep CPU DiffusionGemma fallback fully CPU-masked so a masked GPU host does not re-expose GPU 0 --- studio/backend/core/inference/llama_cpp.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index b07bd33076..f61402aa5c 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3971,7 +3971,11 @@ class LlamaCppBackend: # Auto-size (0): the visual server probes the largest context that fits this GPU's VRAM # (capped at the training context). An explicit in-range n_ctx overrides it. maxtok = n_ctx if (n_ctx and 0 < n_ctx <= 65536) else 0 - gpu = os.environ.get("DG_GPU", "0") + # No visible CUDA GPU: a genuine CPU host, or a GPU host masked with + # CUDA_VISIBLE_DEVICES="" to force CPU serving. Keep the visual-server child + # CPU-masked (empty --gpu) so the shim does not re-expose GPU 0 via its default. + cpu_only = self._effective_gpu_count() == 0 + gpu = "" if cpu_only else os.environ.get("DG_GPU", "0") cmd = list(shim_cmd) + [ "--gguf", @@ -3991,6 +3995,12 @@ class LlamaCppBackend: # refuses to load unless UNSLOTH_IS_PRESENT is set (normally by `import # unsloth`). The shim never imports unsloth, so set it here as unsloth does. env["UNSLOTH_IS_PRESENT"] = "1" + # The shim's `import unsloth_zoo` aborts in get_device_type() ("needs a GPU") + # when no accelerator is visible, even though it only drives the CPU + # visual-server binary and does no torch GPU work. Allow the CPU device so the + # runner starts; the visual server still runs on the CPU llama.cpp build. + if cpu_only: + env.setdefault("UNSLOTH_ALLOW_CPU", "1") env["DG_VISUAL_BIN"] = visual_bin env["DG_GPU"] = gpu # The file-override shim imports its sibling visual_engine; put its dir on PYTHONPATH. From 1a274c488e9621281f86e2f6e85316a77e2c8070 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 8 Jul 2026 07:51:53 -0700 Subject: [PATCH 031/402] Bump install.sh / install.ps1 pins to unsloth>=2026.7.2 and unsloth-zoo>=2026.7.2 (#6981) PyPI release unsloth 2026.7.2 is now live. Bumps the pinned floor in install.sh and install.ps1 from 2026.7.1 to 2026.7.2 for both unsloth and unsloth-zoo across all 5 install commands (no-torch / reinstall / upgrade / local / auto torch backend paths) so fresh installs resolve to the new wheel. Follows the same pattern as #5716. --- install.ps1 | 10 +++++----- install.sh | 10 +++++----- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/install.ps1 b/install.ps1 index 9114f80af9..696f4e613a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -2155,7 +2155,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Resolve pydantic WITH deps so pip pins pydantic-core # to the matching version (no-torch-runtime.txt below @@ -2169,7 +2169,7 @@ exit 0 } } } else { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red @@ -2235,7 +2235,7 @@ exit 0 if ($SkipTorch) { # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } if ($baseInstallExit -eq 0) { # Same pydantic-with-deps trick as the migrated branch. $baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic } @@ -2247,7 +2247,7 @@ exit 0 } } } elseif ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" } } else { $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" } } @@ -2275,7 +2275,7 @@ exit 0 Write-TauriLog "STEP" "Installing unsloth" substep "installing unsloth (this may take a few minutes)..." if ($StudioLocalInstall) { - $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto } + $baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto } if ($baseInstallExit -ne 0) { Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) diff --git a/install.sh b/install.sh index 796d80e401..0acc9ec0be 100755 --- a/install.sh +++ b/install.sh @@ -2706,7 +2706,7 @@ if [ "$_MIGRATED" = true ]; then # to prevent transitive torch resolution. run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Resolve pydantic WITH deps so pip pins pydantic-core to the # matching version (no-torch-runtime.txt below is --no-deps). # All transitive deps are torch-free. @@ -2721,7 +2721,7 @@ if [ "$_MIGRATED" = true ]; then # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ - "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" ${_MLX_LM_EXCLUDE_ARG:-} + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" ${_MLX_LM_EXCLUDE_ARG:-} fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2925,7 +2925,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \ --upgrade-package unsloth --upgrade-package unsloth-zoo \ - "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" + "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" # Same pydantic-with-deps trick as the migrated branch. run_install_cmd_retry "install pydantic (with deps for compatible core)" \ uv pip install --python "$_VENV_PY" pydantic @@ -2943,7 +2943,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ - --upgrade-package unsloth "unsloth>=2026.7.1" "unsloth-zoo>=2026.7.1" + --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." @@ -2975,7 +2975,7 @@ else tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." if [ "$STUDIO_LOCAL_INSTALL" = true ]; then - run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.1" "unsloth>=2026.7.1" --torch-backend=auto + run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps substep "overlaying unsloth-zoo from git main..." From 5c2e53606e513cab3b698852b928b2b4484ada76 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:14:03 -0700 Subject: [PATCH 032/402] Studio: render thinking blocks for safetensors inference with prefilled templates (#6816) * Studio: render thinking blocks for safetensors inference with prefilled templates Reasoning templates like Qwen3.6 end the generation prompt with an open tag. skip_prompt streaming drops it, so the frontend never sees the opening tag and shows reasoning as plain text. Detect the prefill and re-emit it at the start of the stream on the transformers and MLX paths. Also stop stripping think tags in _clean_generated_text when a tokenizer marks them special. * Studio: guard think re-emit for special close tags, yield prefill early Address review feedback: - Guard: skip re-emitting the open when the tokenizer marks as a special token, since skip_special_tokens would strip the model's close tag and leave an unclosed block that swallows the answer. Falls back to plain text (pre-fix behaviour) for those tokenizers. - Yield the prefilled before the first token so the thinking block renders during prompt prefill instead of after the first generated token. - Drop the now-unnecessary _clean_generated_text think-tag exemption; the guard handles the special-token case at the source. No mainstream reasoning model (Qwen3.6, Qwen3, DeepSeek-R1, QwQ, GLM-4.6) marks think tags special, so behaviour is unchanged for them. * [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> Co-authored-by: Lyxot --- .../core/inference/chat_template_helpers.py | 37 ++++++++ studio/backend/core/inference/inference.py | 31 ++++++- .../backend/core/inference/mlx_inference.py | 20 ++++- .../tests/test_think_prefill_reemit.py | 89 +++++++++++++++++++ 4 files changed, 173 insertions(+), 4 deletions(-) create mode 100644 studio/backend/tests/test_think_prefill_reemit.py diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index dfd4c1c0bc..897db8262d 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -12,6 +12,43 @@ import json import logging from typing import Optional +_THINK_OPEN = "" +_THINK_CLOSE = "" + + +def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str: + """Return the trailing open ```` prefill of a rendered prompt. + + Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation + prompt with ``\\n`` so the model starts reasoning immediately. + Because that opening tag is part of the *prompt*, skip_prompt streaming + never emits it, and the frontend's ````/```` parser shows + the reasoning as plain text instead of a thinking block. (The GGUF path + is unaffected: llama-server's reasoning parser returns + ``reasoning_content``, which gets re-wrapped in think tags.) + + Returns the exact prompt tail to re-emit at the start of the generated + stream (e.g. ``"\\n"``), or ``""`` when the prompt does not end + with an open think block, including the ``enable_thinking=False`` case + where templates prefill an already-closed ``\\n\\n``. + + ``special_tokens`` is the tokenizer's special-token list. If ```` + is one, the streamer's skip_special_tokens strips the model's closing tag, + so re-emitting the open would leave an unclosed block that swallows the + answer. In that case return ``""`` and fall back to plain text. + """ + if not prompt: + return "" + open_idx = prompt.rfind(_THINK_OPEN) + if open_idx == -1: + return "" + tail = prompt[open_idx:] + if _THINK_CLOSE in tail or tail.strip() != _THINK_OPEN: + return "" + if special_tokens and _THINK_CLOSE in set(special_tokens): + return "" + return tail + logger = logging.getLogger(__name__) diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 167706f701..7e69e05124 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -1178,13 +1178,22 @@ class InferenceBackend: add_special_tokens = False, return_tensors = "pt", ).to(model.device) + prompt_text = input_text else: # Text-only path for a vision model formatted_prompt = self.format_chat_prompt(messages, system_prompt) inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device) + prompt_text = formatted_prompt # Stream with TextIteratorStreamer + background thread try: + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill swallowed by skip_prompt (see + # generate_stream). + think_prefix = detect_think_prefill( + prompt_text, getattr(raw_tokenizer, "all_special_tokens", None) + ) from transformers import TextIteratorStreamer import threading @@ -1233,7 +1242,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False @@ -1467,6 +1480,16 @@ class InferenceBackend: from transformers import TextIteratorStreamer import threading + from core.inference.chat_template_helpers import detect_think_prefill + + # skip_prompt swallows an open prefilled by the template; + # re-emit it so the frontend can render the thinking block. + # gpt-oss emits its own tags via HarmonyTextStreamer. + think_prefix = ( + "" + if self._is_gpt_oss_model() + else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None)) + ) # gpt-oss models: HarmonyTextStreamer parses the multi-channel # harmony protocol into tags @@ -1550,7 +1573,11 @@ class InferenceBackend: thread = threading.Thread(target = generate_fn) thread.start() - output = "" + output = think_prefix + # Emit the prefilled before the first token so the block + # renders during prompt prefill (which can take seconds). + if think_prefix: + yield think_prefix from queue import Empty generation_complete = False diff --git a/studio/backend/core/inference/mlx_inference.py b/studio/backend/core/inference/mlx_inference.py index d84baa278d..6287b184a6 100644 --- a/studio/backend/core/inference/mlx_inference.py +++ b/studio/backend/core/inference/mlx_inference.py @@ -485,6 +485,7 @@ class MLXInferenceBackend: from core.inference.chat_template_helpers import ( apply_chat_template_for_generation, + detect_think_prefill, render_with_native_template_fallback, ) @@ -518,6 +519,15 @@ class MLXInferenceBackend: hf_token = model_info.get("hf_token"), ) + # An open prefilled by the template lives in the prompt, not + # the generated tokens; re-emit it so the frontend renders the block. + think_prefix = detect_think_prefill( + prompt, getattr(self._tokenizer, "all_special_tokens", None) + ) + # Emit it before the first token so the block renders during prefill. + if think_prefix: + yield think_prefix + sampler = make_sampler( temp = temperature, top_p = top_p, @@ -570,7 +580,7 @@ class MLXInferenceBackend: token_ids, skip_special_tokens = True, ) - yield cumulative + yield think_prefix + cumulative if cancel_event and cancel_event.is_set(): break @@ -634,7 +644,13 @@ class MLXInferenceBackend: # mlx_vlm's stream_generate handles pixel_values (None for text-only) images = [image] if image is not None else None - cumulative = "" + from core.inference.chat_template_helpers import detect_think_prefill + + # Re-emit an open prefill from the prompt (see _generate_text). + cumulative = detect_think_prefill(prompt, getattr(chat_target, "all_special_tokens", None)) + # Emit it before the first token so the block renders during prefill. + if cumulative: + yield cumulative logger.info( "VLM generating: prompt_len=%d, has_image=%s", len(prompt), diff --git a/studio/backend/tests/test_think_prefill_reemit.py b/studio/backend/tests/test_think_prefill_reemit.py new file mode 100644 index 0000000000..300ff92776 --- /dev/null +++ b/studio/backend/tests/test_think_prefill_reemit.py @@ -0,0 +1,89 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +""" +Unit tests for detect_think_prefill. + +Reasoning templates (Qwen3.6-style) end the generation prompt with an open +``\\n`` so the model starts reasoning immediately. skip_prompt +streaming drops that opening tag, so the safetensors/MLX paths must re-emit +it for the frontend's parser to render a thinking block. +""" + +import os +import sys + +_backend = os.path.join(os.path.dirname(__file__), "..") +sys.path.insert(0, _backend) + +from core.inference.chat_template_helpers import detect_think_prefill + + +QWEN_PROMPT = "<|im_start|>user\nHi!<|im_end|>\n<|im_start|>assistant\n" + + +def test_open_think_prefill_reemitted(): + """Qwen3.6-style enable_thinking=True prompt tail: \\n.""" + assert detect_think_prefill(QWEN_PROMPT + "\n") == "\n" + + +def test_bare_open_think_prefill_reemitted(): + """Prefill without trailing newline still detected.""" + assert detect_think_prefill(QWEN_PROMPT + "") == "" + + +def test_closed_think_prefill_not_reemitted(): + """enable_thinking=False prefills a closed, empty think block.""" + assert detect_think_prefill(QWEN_PROMPT + "\n\n\n\n") == "" + + +def test_prompt_without_think_untouched(): + """Non-reasoning templates produce no prefix.""" + assert detect_think_prefill(QWEN_PROMPT) == "" + + +def test_historical_think_blocks_ignored(): + """A closed think block in a prior assistant turn (preserve_thinking) + must not trigger re-emission when the generation tail is plain.""" + prompt = ( + "<|im_start|>user\nHi!<|im_end|>\n" + "<|im_start|>assistant\n\nprior reasoning\n\n\nHello!<|im_end|>\n" + "<|im_start|>user\nAgain?<|im_end|>\n<|im_start|>assistant\n" + ) + assert detect_think_prefill(prompt) == "" + + +def test_historical_blocks_plus_open_prefill(): + """Prior closed blocks plus a fresh open prefill: only the tail matters.""" + prompt = ( + "<|im_start|>assistant\n\nprior\n\n\nHello!<|im_end|>\n" + "<|im_start|>assistant\n\n" + ) + assert detect_think_prefill(prompt) == "\n" + + +def test_content_after_open_tag_not_reemitted(): + """If non-whitespace follows the tag it is not a plain prefill.""" + assert detect_think_prefill(QWEN_PROMPT + "\npartial reasoning") == "" + + +def test_empty_and_none_prompts(): + assert detect_think_prefill("") == "" + assert detect_think_prefill(None) == "" + + +def test_guard_suppresses_when_close_tag_is_special(): + """If is a special token, skip_special_tokens strips the model's + close tag, so re-emitting the open would leave an unclosed block. Guard off.""" + specials = ["<|im_end|>", "", ""] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "" + + +def test_guard_emits_when_think_not_special(): + specials = ["<|im_end|>", "<|endoftext|>"] + assert detect_think_prefill(QWEN_PROMPT + "\n", specials) == "\n" + + +def test_guard_default_and_empty_keep_emitting(): + assert detect_think_prefill(QWEN_PROMPT + "\n", None) == "\n" + assert detect_think_prefill(QWEN_PROMPT + "\n", []) == "\n" From 7a9fb4404e5c81ef5eb34de7d7944bd476c6192b Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Wed, 8 Jul 2026 08:17:09 -0700 Subject: [PATCH 033/402] Remove API menu new badge (#6983) --- studio/frontend/src/components/app-sidebar.tsx | 3 --- 1 file changed, 3 deletions(-) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index f59a952b3a..823e420869 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -1572,9 +1572,6 @@ export function AppSidebar() { > {t("shell.navigation.api")} - - {t("common.new")} - } From 92c3e48529c8b7f96033f52f93845819b2ae53e3 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 8 Jul 2026 13:37:44 -0700 Subject: [PATCH 034/402] Fix BAD_MAPPINGS not redirecting the -unsloth-bnb-4bit dynamic quants (#6949) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .github/workflows/consolidated-tests-ci.yml | 1 + tests/test_bad_mappings_redirect.py | 47 +++++++++++++++++++++ unsloth/models/loader_utils.py | 5 +++ 3 files changed, 53 insertions(+) create mode 100644 tests/test_bad_mappings_redirect.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index ae4b386589..6ff3d19ba2 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -364,6 +364,7 @@ jobs: tests/utils/test_attention_masks.py \ tests/utils/test_trunc_normal_patch.py \ tests/python/test_fast_language_model_text_only.py \ + tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is diff --git a/tests/test_bad_mappings_redirect.py b/tests/test_bad_mappings_redirect.py new file mode 100644 index 0000000000..49dab2d98b --- /dev/null +++ b/tests/test_bad_mappings_redirect.py @@ -0,0 +1,47 @@ +"""Regression test for BAD_MAPPINGS redirecting oversized dynamic quants. + +get_model_name previously applied BAD_MAPPINGS only to the resolver's output, +but several listed names (the `-unsloth-bnb-4bit` dynamic quants, plus any name +the resolver doesn't map) come back as None, so their BAD_MAPPINGS entries were +dead and the oversized model loaded. Asserting over every entry catches all of +them. The mapper table and the resolver have no heavy imports of their own, +so we exec the import-free mapper module and ast-extract the resolver functions +rather than importing unsloth (which needs a GPU). +""" + +import ast +import os + +_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models") + + +def _load_get_model_name(): + mapper_ns = {} + with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f: + exec(compile(f.read(), "mapper.py", "exec"), mapper_ns) + + with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f: + tree = ast.parse(f.read()) + + namespace = dict(mapper_ns) + namespace["SUPPORTS_FOURBIT"] = True + namespace["_env_says_offline"] = lambda: True + namespace["_get_new_mapper"] = lambda: ({}, {}, {}) + + wanted = {"__get_model_name", "_resolve_with_mappers", "get_model_name"} + for node in tree.body: + if isinstance(node, ast.Assign) and any( + getattr(target, "id", None) == "BAD_MAPPINGS" for target in node.targets + ): + exec(compile(ast.Module([node], []), "", "exec"), namespace) + elif isinstance(node, ast.FunctionDef) and node.name in wanted: + exec(compile(ast.Module([node], []), node.name, "exec"), namespace) + + return namespace["get_model_name"], namespace["BAD_MAPPINGS"] + + +def test_bad_mappings_redirect_every_listed_name(): + get_model_name, bad_mappings = _load_get_model_name() + assert bad_mappings, "BAD_MAPPINGS should not be empty" + for name, expected in bad_mappings.items(): + assert get_model_name(name, load_in_4bit = True) == expected, name diff --git a/unsloth/models/loader_utils.py b/unsloth/models/loader_utils.py index d6d6ce877e..fa6282bcf6 100644 --- a/unsloth/models/loader_utils.py +++ b/unsloth/models/loader_utils.py @@ -239,6 +239,11 @@ def get_model_name( and new_model_name.lower() in BAD_MAPPINGS ): new_model_name = BAD_MAPPINGS[new_model_name.lower()] + elif new_model_name is None and model_name.lower() in BAD_MAPPINGS: + # Some bad names (e.g. the `-unsloth-bnb-4bit` dynamic quants) are keys + # of the mappers, not values, so the resolver returns None for them and + # the remap above is skipped; remap the input name directly instead. + new_model_name = BAD_MAPPINGS[model_name.lower()] if ( new_model_name is None From dc4618ce475554d99fab76daae9746baa62bf090 Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 8 Jul 2026 13:40:39 -0700 Subject: [PATCH 035/402] Fix duplicate unsloth/gemma-2b-bnb-4bit mapper key routing the base 4bit repo to the instruct model (#6891) --- .github/workflows/consolidated-tests-ci.yml | 1 + tests/test_gemma_2b_mapper_key.py | 46 +++++++++++++++++++++ unsloth/models/mapper.py | 2 +- 3 files changed, 48 insertions(+), 1 deletion(-) create mode 100644 tests/test_gemma_2b_mapper_key.py diff --git a/.github/workflows/consolidated-tests-ci.yml b/.github/workflows/consolidated-tests-ci.yml index 6ff3d19ba2..1bb4c2bb58 100644 --- a/.github/workflows/consolidated-tests-ci.yml +++ b/.github/workflows/consolidated-tests-ci.yml @@ -366,6 +366,7 @@ jobs: tests/python/test_fast_language_model_text_only.py \ tests/test_bad_mappings_redirect.py \ tests/test_prefetch_snapshot_scope.py \ + tests/test_gemma_2b_mapper_key.py \ --deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap' # The deselected test monkeypatches flash_attn_varlen_func, which is # only bound on the module when `flash_attn` is importable. flash_attn diff --git a/tests/test_gemma_2b_mapper_key.py b/tests/test_gemma_2b_mapper_key.py new file mode 100644 index 0000000000..31edacfce4 --- /dev/null +++ b/tests/test_gemma_2b_mapper_key.py @@ -0,0 +1,46 @@ +"""Regression test for the duplicate ``unsloth/gemma-2b-bnb-4bit`` key in +``unsloth/models/mapper.py``. + +The 4bit instruction-tuned Gemma 2B entry was accidentally keyed with the base +model's repo name, so ``__INT_TO_FLOAT_MAPPER`` held two identical +``unsloth/gemma-2b-bnb-4bit`` keys. Python keeps only the last value for a +duplicate literal key, so the base 4bit repo resolved to the *instruct* model, +the base model lost its reverse (4x-faster) mapping, and +``unsloth/gemma-2b-it-bnb-4bit`` was never registered at all. + +``mapper.py`` has no imports, so we exec it directly and inspect the built +mappers without importing ``unsloth`` (which requires a GPU). +""" + +import os + +MAPPER_PATH = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models", "mapper.py") + + +def _load_mappers(): + with open(MAPPER_PATH) as f: + source = f.read() + namespace = {} + exec(compile(source, MAPPER_PATH, "exec"), namespace) + return namespace + + +def test_gemma_2b_base_and_instruct_4bit_are_distinct(): + namespace = _load_mappers() + int_to_float = namespace["INT_TO_FLOAT_MAPPER"] + float_to_int = namespace["FLOAT_TO_INT_MAPPER"] + + # The base 4bit repo must resolve to the base model, not the instruct one. + assert int_to_float["unsloth/gemma-2b-bnb-4bit"] == "unsloth/gemma-2b" + + # The instruct 4bit repo must be registered and resolve to the instruct model. + assert "unsloth/gemma-2b-it-bnb-4bit" in int_to_float + assert int_to_float["unsloth/gemma-2b-it-bnb-4bit"] == "unsloth/gemma-2b-it" + + # The base model must reverse-map back to the base 4bit repo. + assert float_to_int["unsloth/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit" + assert float_to_int["google/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit" + + # The instruct model must reverse-map to the instruct 4bit repo. + assert float_to_int["unsloth/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit" + assert float_to_int["google/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit" diff --git a/unsloth/models/mapper.py b/unsloth/models/mapper.py index 57c1e292c3..f3a0e1f9bb 100644 --- a/unsloth/models/mapper.py +++ b/unsloth/models/mapper.py @@ -134,7 +134,7 @@ __INT_TO_FLOAT_MAPPER = \ "unsloth/gemma-7b-it", "google/gemma-7b-it", ), - "unsloth/gemma-2b-bnb-4bit" : ( + "unsloth/gemma-2b-it-bnb-4bit" : ( "unsloth/gemma-2b-it", "google/gemma-2b-it", ), From 85a068cfe10f2f8bc214dbd2e0b82164419fc8de Mon Sep 17 00:00:00 2001 From: Vineeth Sai Date: Wed, 8 Jul 2026 13:57:41 -0700 Subject: [PATCH 036/402] Fix to_sharegpt optional block rendering "None" for missing extra columns (#6827) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- .../python/test_to_sharegpt_optional_none.py | 97 +++++++++++++++++++ unsloth/chat_templates.py | 13 ++- 2 files changed, 108 insertions(+), 2 deletions(-) create mode 100644 tests/python/test_to_sharegpt_optional_none.py diff --git a/tests/python/test_to_sharegpt_optional_none.py b/tests/python/test_to_sharegpt_optional_none.py new file mode 100644 index 0000000000..05fcb8a6a8 --- /dev/null +++ b/tests/python/test_to_sharegpt_optional_none.py @@ -0,0 +1,97 @@ +import ast +import re +from pathlib import Path + + +def _load_formatter_builders(): + # Extract _parse_combined_prompt and _create_formatter without importing + # unsloth (importing unsloth needs unsloth_zoo / a GPU). Both are pure + # Python and only use the `re` module. + source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py" + tree = ast.parse(source.read_text(encoding = "utf-8")) + wanted = {"_parse_combined_prompt", "_create_formatter"} + funcs = [ + node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted + ] + namespace = {"re": re} + module = ast.Module(body = funcs, type_ignores = []) + ast.fix_missing_locations(module) + exec(compile(module, str(source), "exec"), namespace) + return namespace["_parse_combined_prompt"], namespace["_create_formatter"] + + +class _StubDataset: + def __init__(self, column_names): + self.column_names = column_names + + +def _render(merged_prompt, columns, batch): + parse, create = _load_formatter_builders() + possible_columns, final_optional_prompts = parse(merged_prompt, _StubDataset(columns)) + processor = create(possible_columns, final_optional_prompts, "text") + return processor(batch)["text"] + + +def test_optional_block_missing_second_column_does_not_render_none(): + # A [[...]] block may reference several columns; only the first gates the + # block. A later column that is None must not render as the literal "None". + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": [None]}, + ) + assert out[0] == "Location: Paris, end" + assert "None" not in out[0] + + +def test_optional_block_all_columns_present_unchanged(): + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": ["France"]}, + ) + assert out[0] == "Location: Paris, France end" + + +def test_optional_block_gating_column_empty_is_dropped(): + # When the gating (first) column is empty the whole block is omitted; this + # behaviour is unchanged by the None coercion. + merged_prompt = "Location: [[{city}, {country}]] end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": [""], "country": ["France"]}, + ) + assert out[0] == "Location: end" + + +def test_single_column_optional_block_gated_out_on_none(): + # Single-column blocks were already gated correctly (the sole column is the + # gate); confirm they stay unaffected. + merged_prompt = "Name: [[{name}]]!" + out = _render(merged_prompt, ["name"], {"name": [None, "Bob"]}) + assert out == ["Name: !", "Name: Bob!"] + + +def test_required_column_none_does_not_render_none(): + # A required (non-[[...]]) column that is None must not render as the + # literal "None" either; coercion happens at the row source, so both the + # required and optional branches are covered. + merged_prompt = "Location: {city}, {country} end" + out = _render( + merged_prompt, + ["city", "country"], + {"city": ["Paris"], "country": [None]}, + ) + assert out[0] == "Location: Paris, end" + assert "None" not in out[0] + + +def test_optional_block_falsy_but_present_gating_value_still_renders(): + # The gate keeps a block whenever the first column is not "". A falsy but + # real value (0) must not be treated as absent, so the block still renders. + merged_prompt = "Count: [[{n}]]!" + out = _render(merged_prompt, ["n"], {"n": [0]}) + assert out[0] == "Count: 0!" diff --git a/unsloth/chat_templates.py b/unsloth/chat_templates.py index 169b2dbd0e..dd1e433471 100644 --- a/unsloth/chat_templates.py +++ b/unsloth/chat_templates.py @@ -2185,7 +2185,16 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name texts = [] for row_idx in range(n_rows): - row_values = {column: examples[column][row_idx] for column in columns} + # Coerce missing (None) columns to "" so they do not render as the + # literal string "None" in the emitted text. In a [[...]] block only + # the first column gates the block, so a later column can still be + # None here; required columns can be None too. Coercing at the source + # covers both; since None is now "", the gate below only needs to + # test for "" (an empty first column still drops the block). + row_values = { + column: ("" if (value := examples[column][row_idx]) is None else value) + for column in columns + } formatter_values = {} for formatter_template in formatter_templates: @@ -2196,7 +2205,7 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name continue _, optional_name, prompt, needed_columns = formatter_template - if row_values[needed_columns[0]] not in (None, ""): + if row_values[needed_columns[0]] != "": prompt_values = {column: row_values[column] for column in needed_columns} formatter_values[optional_name] = prompt.format(**prompt_values) else: From 81f789ba85bb45c30fe7a8e60126112e07e3ddac Mon Sep 17 00:00:00 2001 From: ramisworld Date: Thu, 9 Jul 2026 09:32:51 +1200 Subject: [PATCH 037/402] Guard FP8 Triton launches with tensor device context (#6888) --------- Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com> --- tests/test_fp8_device_context.py | 249 +++++++++++++++++++++++++++++++ unsloth/kernels/fp8.py | 87 ++++++----- 2 files changed, 300 insertions(+), 36 deletions(-) create mode 100644 tests/test_fp8_device_context.py diff --git a/tests/test_fp8_device_context.py b/tests/test_fp8_device_context.py new file mode 100644 index 0000000000..2eea35f4e6 --- /dev/null +++ b/tests/test_fp8_device_context.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +import ast +from contextlib import nullcontext +from pathlib import Path +from types import SimpleNamespace + +import pytest + + +REPO_ROOT = Path(__file__).resolve().parents[1] +FP8_SOURCE = REPO_ROOT / "unsloth" / "kernels" / "fp8.py" + + +class _FakeDeviceModule: + def __init__(self, device_count: int) -> None: + self._device_count = device_count + self.device_calls = [] + + def device_count(self) -> int: + return self._device_count + + def device(self, device): + self.device_calls.append(device) + return ("device-context", device) + + +class _FakeTorch: + Tensor = object + + def __init__( + self, + cuda_device_count: int, + xpu_device_count: int = 0, + ) -> None: + self.cuda = _FakeDeviceModule(cuda_device_count) + self.xpu = _FakeDeviceModule(xpu_device_count) + + +class _LaunchVisitor(ast.NodeVisitor): + def __init__(self) -> None: + self.guarded_launches: set[str] = set() + self.unguarded_launches: set[str] = set() + self._inside_fp8_device_context = 0 + + def visit_With(self, node: ast.With) -> None: + enters_context = any( + isinstance(item.context_expr, ast.Call) + and isinstance(item.context_expr.func, ast.Name) + and item.context_expr.func.id == "_fp8_triton_device_context" + for item in node.items + ) + if enters_context: + self._inside_fp8_device_context += 1 + for statement in node.body: + self.visit(statement) + if enters_context: + self._inside_fp8_device_context -= 1 + + def visit_Call(self, node: ast.Call) -> None: + launch_name = self._triton_launch_name(node) + if launch_name is not None: + if self._inside_fp8_device_context: + self.guarded_launches.add(launch_name) + else: + self.unguarded_launches.add(launch_name) + self.generic_visit(node) + + @staticmethod + def _triton_launch_name(node: ast.Call) -> str | None: + if isinstance(node.func, ast.Name) and node.func.id == "triton_quantize_fp8_block": + return node.func.id + if not isinstance(node.func, ast.Subscript): + return None + if not isinstance(node.func.value, ast.Name): + return None + return node.func.value.id + + +def _load_device_context_helper(fake_torch: _FakeTorch): + source = FP8_SOURCE.read_text() + tree = ast.parse(source) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == "_fp8_triton_device_context": + namespace = {"torch": fake_torch, "nullcontext": nullcontext} + exec(ast.get_source_segment(source, node), namespace) + return namespace["_fp8_triton_device_context"] + raise AssertionError("_fp8_triton_device_context was not found") + + +def test_fp8_device_context_selects_cuda_tensor_device_on_multi_gpu() -> None: + fake_torch = _FakeTorch(cuda_device_count = 2) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda")) + + context = helper(tensor) + + assert context == ("device-context", tensor.device) + assert fake_torch.cuda.device_calls == [tensor.device] + + +def test_fp8_device_context_is_noop_for_single_cuda_device() -> None: + fake_torch = _FakeTorch(cuda_device_count = 1) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda")) + + context = helper(tensor) + + assert isinstance(context, nullcontext) + assert fake_torch.cuda.device_calls == [] + + +def test_fp8_device_context_selects_xpu_tensor_device_on_multi_gpu() -> None: + fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 2) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu")) + + context = helper(tensor) + + assert context == ("device-context", tensor.device) + assert fake_torch.xpu.device_calls == [tensor.device] + + +def test_fp8_device_context_is_noop_for_single_xpu_device() -> None: + fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 1) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu")) + + context = helper(tensor) + + assert isinstance(context, nullcontext) + assert fake_torch.xpu.device_calls == [] + + +def test_fp8_device_context_is_noop_for_non_cuda_tensor() -> None: + fake_torch = _FakeTorch(cuda_device_count = 8) + helper = _load_device_context_helper(fake_torch) + tensor = SimpleNamespace(device = SimpleNamespace(type = "cpu")) + + context = helper(tensor) + + assert isinstance(context, nullcontext) + assert fake_torch.cuda.device_calls == [] + + +def test_fp8_triton_launches_enter_tensor_device_context() -> None: + tree = ast.parse(FP8_SOURCE.read_text()) + function_names = {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)} + assert "_fp8_triton_device_context" in function_names + + visitor = _LaunchVisitor() + visitor.visit(tree) + + expected_launches = { + "weight_dequant_kernel", + "act_quant_kernel", + "_w8a8_block_fp8_matmul", + "triton_quantize_fp8_block", + } + assert expected_launches <= visitor.guarded_launches + assert not (expected_launches & visitor.unguarded_launches) + + +def _require_two_cuda_devices(): + torch = pytest.importorskip("torch") + pytest.importorskip("triton") + + if not torch.cuda.is_available() or torch.cuda.device_count() < 2: + pytest.skip("requires at least two CUDA devices") + return torch + + +def test_weight_dequant_block_runs_on_tensor_device_when_current_device_differs() -> None: + torch = _require_two_cuda_devices() + from unsloth.kernels.fp8 import weight_dequant_block + + previous_device = torch.cuda.current_device() + try: + torch.cuda.set_device(0) + x = torch.arange(256 * 256, device = "cuda:1", dtype = torch.float32).reshape(256, 256) + scales = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device = "cuda:1", dtype = torch.float32) + + actual = weight_dequant_block(x, scales, block_size = 128, dtype = torch.float32) + + expanded_scales = scales.repeat_interleave(128, dim = 0).repeat_interleave(128, dim = 1) + expected = x * expanded_scales + + assert actual.device == x.device + assert torch.cuda.current_device() == 0 + torch.testing.assert_close(actual, expected) + finally: + torch.cuda.set_device(previous_device) + + +def test_act_quant_runs_on_tensor_device_when_current_device_differs() -> None: + torch = _require_two_cuda_devices() + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("requires torch.float8_e4m3fn") + if torch.cuda.get_device_capability(1)[0] < 9: + pytest.skip("requires FP8-capable CUDA hardware") + + from unsloth.kernels.fp8 import act_quant + + previous_device = torch.cuda.current_device() + try: + torch.cuda.set_device(0) + x = torch.arange(256, device = "cuda:1", dtype = torch.float32).reshape(2, 128) + + y, scales = act_quant(x, block_size = 128) + + assert y.device == x.device + assert scales.device == x.device + assert torch.cuda.current_device() == 0 + finally: + torch.cuda.set_device(previous_device) + + +def test_w8a8_block_fp8_matmul_triton_runs_on_tensor_device_when_current_device_differs() -> None: + torch = _require_two_cuda_devices() + if not hasattr(torch, "float8_e4m3fn"): + pytest.skip("requires torch.float8_e4m3fn") + if torch.cuda.get_device_capability(1)[0] < 9: + pytest.skip("requires FP8-capable CUDA hardware") + + from unsloth.kernels.fp8 import w8a8_block_fp8_matmul_triton + + previous_device = torch.cuda.current_device() + try: + torch.cuda.set_device(0) + A = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn) + B = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn) + As = torch.ones((128, 1), device = "cuda:1", dtype = torch.float32) + Bs = torch.ones((1, 1), device = "cuda:1", dtype = torch.float32) + + actual = w8a8_block_fp8_matmul_triton( + A, + B, + As, + Bs, + block_size = [128, 128], + output_dtype = torch.float32, + ) + + expected = torch.full((128, 128), 128.0, device = "cuda:1", dtype = torch.float32) + assert actual.device == A.device + assert torch.cuda.current_device() == 0 + torch.testing.assert_close(actual, expected) + finally: + torch.cuda.set_device(previous_device) diff --git a/unsloth/kernels/fp8.py b/unsloth/kernels/fp8.py index 935ffbb447..4efc4bd5d3 100644 --- a/unsloth/kernels/fp8.py +++ b/unsloth/kernels/fp8.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import os +from contextlib import nullcontext import torch import torch.nn as nn import triton @@ -24,6 +25,15 @@ from unsloth_zoo.temporary_patches.common import torch_compile torch_matmul = torch.matmul + +def _fp8_triton_device_context(tensor: torch.Tensor): + if tensor.device.type == "cuda" and torch.cuda.device_count() > 1: + return torch.cuda.device(tensor.device) + if tensor.device.type == "xpu" and hasattr(torch, "xpu") and torch.xpu.device_count() > 1: + return torch.xpu.device(tensor.device) + return nullcontext() + + try: from transformers.integrations.finegrained_fp8 import FP8Linear except: @@ -95,7 +105,8 @@ def weight_dequant_block( triton.cdiv(M, meta["BLOCK_SIZE"]), triton.cdiv(N, meta["BLOCK_SIZE"]), ) - weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size) + with _fp8_triton_device_context(x): + weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size) return y @@ -149,7 +160,8 @@ def act_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, tor def grid(meta): return (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),) - act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size) + with _fp8_triton_device_context(x): + act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size) return y, s @@ -274,32 +286,33 @@ def w8a8_block_fp8_matmul_triton( def grid(META): return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),) - _w8a8_block_fp8_matmul[grid]( - A, - B, - C, - As, - Bs, - M, - N, - K, - block_n, - block_k, - A.stride(-2), - A.stride(-1), - B.stride(1), - B.stride(0), - C.stride(-2), - C.stride(-1), - As.stride(-2), - As.stride(-1), - Bs.stride(1), - Bs.stride(0), - BLOCK_SIZE_M = BLOCK_SIZE_M, - BLOCK_SIZE_N = BLOCK_SIZE_N, - BLOCK_SIZE_K = BLOCK_SIZE_K, - GROUP_SIZE_M = 8, - ) + with _fp8_triton_device_context(A): + _w8a8_block_fp8_matmul[grid]( + A, + B, + C, + As, + Bs, + M, + N, + K, + block_n, + block_k, + A.stride(-2), + A.stride(-1), + B.stride(1), + B.stride(0), + C.stride(-2), + C.stride(-1), + As.stride(-2), + As.stride(-1), + Bs.stride(1), + Bs.stride(0), + BLOCK_SIZE_M = BLOCK_SIZE_M, + BLOCK_SIZE_N = BLOCK_SIZE_N, + BLOCK_SIZE_K = BLOCK_SIZE_K, + GROUP_SIZE_M = 8, + ) return C @@ -311,13 +324,14 @@ def torchao_block_matmul( block_size: tuple[int, int], output_dtype: torch.dtype = torch.bfloat16, ): - out = torchao_blockwise_gemm( - act_q.contiguous(), - act_scale.contiguous(), - weight_q.contiguous(), - weight_scale.contiguous(), - block_size = block_size[1], - ) + with _fp8_triton_device_context(act_q): + out = torchao_blockwise_gemm( + act_q.contiguous(), + act_scale.contiguous(), + weight_q.contiguous(), + weight_scale.contiguous(), + block_size = block_size[1], + ) return out.to(output_dtype) @@ -540,7 +554,8 @@ class FP8_fbgemm_block_linear(torch.autograd.Function): f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {bs_n, bs_k}" ) - xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None) + with _fp8_triton_device_context(X): + xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None) # TODO: WARNING - diverges from baseline for high X values, producing # gibberish / high starting loss. Do not use until resolved; kept for a # future headstart. From 3b73cd88293728b3ba173e235d405cec40f7a42b Mon Sep 17 00:00:00 2001 From: Nilay <118994073+NilayYadav@users.noreply.github.com> Date: Thu, 9 Jul 2026 08:33:03 +0530 Subject: [PATCH 038/402] Fix per-block ID collisions and add block cleanup for unstructured uploads (#6944) * unstructured block removal * Enhance unstructured block handling * Restrict block cleanup to upload UIDs * cleanup for seed block uploads * upload cleanup queue for unstructured blocks in recipe studio * Fix unstructured upload cleanup edge cases * Fix unstructured upload import ownership * Fix-unstructured-import-path-ownership * Guard failed-delete restore against stale block in unstructured drop zone * Drain queued upload cleanups when autosave is skipped --------- Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: imagineer99 Co-authored-by: Daniel Han --- studio/backend/routes/data_recipe/seed.py | 37 ++++++ studio/backend/tests/test_data_recipe_seed.py | 97 +++++++++++++++ .../src/features/recipe-studio/api/index.ts | 10 ++ .../dialogs/seed/seed-dialog.tsx | 69 ++++++++++- .../dialogs/seed/unstructured-drop-zone.tsx | 37 ++++-- .../hooks/use-recipe-persistence.ts | 86 ++++++++++++-- .../recipe-studio/stores/recipe-studio.ts | 61 +++++++++- .../src/features/recipe-studio/types/index.ts | 2 + .../recipe-studio/utils/config-factories.ts | 43 +++++++ .../recipe-studio/utils/import/importer.ts | 112 +++++++++++------- .../import/parsers/seed-config-parser.ts | 14 ++- .../utils/payload/build-payload.ts | 3 + .../recipe-studio/utils/payload/types.ts | 2 + 13 files changed, 502 insertions(+), 71 deletions(-) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 57a291291e..a5b75b7335 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -10,6 +10,7 @@ import binascii import json import os import re +import shutil from itertools import islice from pathlib import Path from typing import Any @@ -59,6 +60,9 @@ UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"} SEED_UPLOAD_DIR = seed_uploads_root() UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root() _SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$") +# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...) +# never match: those directories can be shared by several recipes. +_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$") def _validate_safe_id(value: str, label: str) -> str: @@ -580,6 +584,39 @@ async def remove_unstructured_file(block_id: str, file_id: str): return {"status": "ok"} +@router.delete("/seed/unstructured-block/{block_id}") +async def remove_unstructured_block(block_id: str): + """Delete a block's upload directory; files on disk still count toward its quota. + + Only uid-namespaced directories may be bulk-deleted: they have exactly one + owning block. Legacy node-id directories (n1, ...) can be shared by other + recipes, so they are managed file-by-file instead. + """ + _validate_safe_id(block_id, "block_id") + if not _UPLOAD_UID_RE.match(block_id): + raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted") + + block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve() + if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()): + raise HTTPException(400, "Invalid block_id: outside upload root") + if not block_dir.exists(): + return {"status": "ok", "deleted": False} + + try: + shutil.rmtree(block_dir) + except OSError as exc: + raise log_and_http_error( + exc, + 500, + "failed to delete uploaded files", + event = "data_recipe.seed.unstructured_block_delete_failed", + log = logger, + ) from exc + if block_dir.exists(): + raise HTTPException(500, "failed to delete uploaded files") + return {"status": "ok", "deleted": True} + + @router.post("/seed/inspect-upload", response_model = SeedInspectResponse) def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse: if payload.file_ids is not None: diff --git a/studio/backend/tests/test_data_recipe_seed.py b/studio/backend/tests/test_data_recipe_seed.py index 09e22116ed..58bbd24061 100644 --- a/studio/backend/tests/test_data_recipe_seed.py +++ b/studio/backend/tests/test_data_recipe_seed.py @@ -124,3 +124,100 @@ def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, e assert result.status == "error" assert result.error == "Text extraction failed." assert _block_files(seed_route) == [] + + +_TEST_UPLOAD_UID = "0f" * 16 + + +def test_remove_unstructured_block_deletes_directory(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = _TEST_UPLOAD_UID) + assert _block_files(seed_route, _TEST_UPLOAD_UID) != [] + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": True} + assert not (seed_route.UNSTRUCTURED_UPLOAD_ROOT / _TEST_UPLOAD_UID).exists() + + +def test_remove_unstructured_block_missing_directory_is_ok(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert result == {"status": "ok", "deleted": False} + + +def test_remove_unstructured_block_rejects_unsafe_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("../escape")) + + assert exc.value.status_code == 400 + + +def test_remove_unstructured_block_rejects_legacy_node_ids(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + _run_upload(seed_route, "notes.txt", b"hello", block_id = "n1") + assert _block_files(seed_route, "n1") != [] + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block("n1")) + + assert exc.value.status_code == 400 + assert _block_files(seed_route, "n1") != [] + + +def test_remove_unstructured_block_rejects_symlink_escape(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + outside = tmp_path / "outside" + outside.mkdir() + (outside / "victim.txt").write_text("keep me") + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + root.mkdir(parents = True) + (root / _TEST_UPLOAD_UID).symlink_to(outside) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert exc.value.status_code == 400 + assert (outside / "victim.txt").exists() + + +def test_remove_unstructured_block_fails_if_directory_remains(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + root = seed_route.UNSTRUCTURED_UPLOAD_ROOT + block_dir = root / _TEST_UPLOAD_UID + block_dir.mkdir(parents = True) + (block_dir / "victim.txt").write_text("keep me") + + calls = [] + + def noop_rmtree(path, *args, **kwargs): + calls.append((path, args, kwargs)) + + monkeypatch.setattr(seed_route.shutil, "rmtree", noop_rmtree) + + with pytest.raises(seed_route.HTTPException) as exc: + asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID)) + + assert calls + assert exc.value.status_code == 500 + assert block_dir.exists() + + +def test_total_upload_quota_is_scoped_per_block(monkeypatch, tmp_path): + seed_route = _load_seed_route(monkeypatch, tmp_path) + monkeypatch.setattr(seed_route, "UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES", 10) + + first = _run_upload(seed_route, "a.txt", b"123456789") + assert first.status == "ok" + + with pytest.raises(seed_route.HTTPException) as exc: + _run_upload(seed_route, "b.txt", b"123") + assert exc.value.status_code == 413 + + # Another block starts with its own untouched budget. + other = _run_upload(seed_route, "c.txt", b"123", block_id = "other") + assert other.status == "ok" diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index 4fe24b1f6d..f4e8167cb3 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -494,3 +494,13 @@ export async function removeUnstructuredFile( throw new Error("Failed to remove file"); } } + +export async function removeUnstructuredBlock(blockId: string): Promise { + const res = await authFetch( + `${DATA_DESIGNER_API_BASE}/seed/unstructured-block/${encodeURIComponent(blockId)}`, + { method: "DELETE" }, + ); + if (!res.ok && res.status !== 404) { + throw new Error("Failed to remove uploaded files"); + } +} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 0ed7eeed75..53f8566090 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -53,6 +53,11 @@ import { inspectSeedDataset, inspectSeedUpload, } from "../../api"; +import { useRecipeStudioStore } from "../../stores/recipe-studio"; +import { + makeUnstructuredUploadUid, + resolveUnstructuredUploadBlockId, +} from "../../utils/config-factories"; import { resolveImagePreview } from "../../utils/image-preview"; import type { GithubItemType, @@ -597,6 +602,41 @@ export function SeedDialog({ const mode = config.seed_source_type ?? "hf"; const previewEmpty = getPreviewEmptyStateCopy(mode); + const queueUploadCleanup = useRecipeStudioStore( + (state) => state.queueUploadCleanup, + ); + + // config.id collides across recipes (ids reset to n1 on import); use a + // stable per-block uid instead. Generate one synchronously so the first + // rendered drop zone cannot upload under a legacy node id. + const uploadUid = config.unstructured_upload_uid?.trim() ?? ""; + const unstructuredFileCount = config.unstructured_file_ids?.length ?? 0; + const generatedUploadUidRef = useRef(null); + if ( + mode === "unstructured" && + !uploadUid && + unstructuredFileCount === 0 && + generatedUploadUidRef.current === null + ) { + generatedUploadUidRef.current = makeUnstructuredUploadUid(); + } + const uploadBlockId = resolveUnstructuredUploadBlockId({ + configId: config.id, + uploadUid, + generatedUploadUid: generatedUploadUidRef.current, + unstructuredFileCount, + }); + + useEffect(() => { + if (mode !== "unstructured") return; + if (uploadUid) return; + if (unstructuredFileCount > 0) return; + const nextUid = + generatedUploadUidRef.current ?? makeUnstructuredUploadUid(); + generatedUploadUidRef.current = nextUid; + onUpdate({ unstructured_upload_uid: nextUid }); + }, [mode, uploadUid, unstructuredFileCount, onUpdate]); + const prevModeRef = useRef(mode); useEffect(() => { const prevMode = prevModeRef.current; @@ -720,6 +760,11 @@ export function SeedDialog({ subset: config.hf_subset?.trim() || undefined, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -730,6 +775,7 @@ export function SeedDialog({ hf_split: response.split ?? "", hf_subset: response.subset ?? "", local_file_name: "", + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -754,6 +800,11 @@ export function SeedDialog({ content_base64: payload, preview_size: 10, }); + // Queue the block's upload directory for deletion after the next + // save; only uid-namespaced directories qualify (single owner). + if (uploadUid && unstructuredFileCount > 0) { + queueUploadCleanup(uploadUid); + } onUpdate({ hf_path: response.resolved_path, seed_columns: response.columns, @@ -765,6 +816,7 @@ export function SeedDialog({ hf_subset: "", hf_split: "", local_file_name: localFile.name, + unstructured_upload_uid: "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -789,7 +841,7 @@ export function SeedDialog({ const { chunkSize, chunkOverlap } = resolveChunking(config); const response = await inspectSeedUpload({ - block_id: config.id, + block_id: uploadBlockId, file_ids: fileIds, file_names: fileNames, preview_size: 10, @@ -827,7 +879,18 @@ export function SeedDialog({ setIsInspecting(false); } }, - [config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles], + [ + config, + getCurrentLoadKey, + localFile, + mode, + onUpdate, + queueUploadCleanup, + unstructuredFiles, + unstructuredFileCount, + uploadBlockId, + uploadUid, + ], ); useEffect(() => { @@ -997,7 +1060,7 @@ export function SeedDialog({ {mode === "unstructured" && ( (null); const filesRef = useRef(files); + const blockIdRef = useRef(blockId); + const mountedRef = useRef(true); const [isDragOver, setIsDragOver] = useState(false); useEffect(() => { filesRef.current = files; - }, [files]); + blockIdRef.current = blockId; + }, [files, blockId]); + useEffect(() => () => { + mountedRef.current = false; + }, []); const totalSize = files.reduce((sum, f) => sum + f.size, 0); @@ -134,15 +140,32 @@ export function UnstructuredDropZone({ if (entry.status === "uploading" && entry.abortController) { entry.abortController.abort(); } - if ( + const needsServerRemove = entry.id && entry.status === "ok" && - !deletedIdsRef.current.has(entry.id) - ) { - deletedIdsRef.current.add(entry.id); - void removeUnstructuredFile(blockId, entry.id).catch(() => {}); - } + !deletedIdsRef.current.has(entry.id); onFilesChange((prev) => prev.filter((_, i) => i !== index)); + if (!needsServerRemove) return; + deletedIdsRef.current.add(entry.id); + removeUnstructuredFile(blockId, entry.id).catch(() => { + // Skip if the drop zone unmounted or its block changed: the id no + // longer belongs here and restoring would leak it into another block. + if (!mountedRef.current || blockIdRef.current !== blockId) return; + // Still exists server-side (counts toward quota); restore it at its + // original position. + deletedIdsRef.current.delete(entry.id); + onFilesChange((prev) => { + const next = [...prev]; + next.splice(Math.min(index, next.length), 0, { + id: entry.id, + name: entry.name, + size: entry.size, + status: "ok", + error: "Remove failed — try again", + }); + return next; + }); + }); }, [blockId, onFilesChange], ); diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts index 0417a6dd22..2d91272469 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts @@ -4,11 +4,13 @@ import { useCallback, useEffect, useMemo, useState } from "react"; import { toastError, toastSuccess } from "@/shared/toast"; import { normalizeNonEmptyName } from "@/utils"; +import { removeUnstructuredBlock } from "../api"; import { buildSignature, copyTextToClipboard, formatSavedLabel, } from "../executions/execution-helpers"; +import { useRecipeStudioStore } from "../stores/recipe-studio"; import { importRecipePayload, type RecipeSnapshot } from "../utils/import"; import type { RecipePayloadResult } from "../utils/payload/types"; @@ -72,7 +74,10 @@ function stripApiKeys(value: unknown): unknown { !Array.isArray(output.env) ) { output.env = Object.fromEntries( - Object.keys(output.env as Record).map((envKey) => [envKey, ""]), + Object.keys(output.env as Record).map((envKey) => [ + envKey, + "", + ]), ); } return output; @@ -82,10 +87,7 @@ function inferHfRepoIdFromPath(pathValue: unknown): string { if (typeof pathValue !== "string") { return ""; } - const parts = pathValue - .trim() - .split("/") - .filter(Boolean); + const parts = pathValue.trim().split("/").filter(Boolean); if (parts.length >= 3 && parts[0] === "datasets") { return `${parts[1]}/${parts[2]}`; } @@ -126,8 +128,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { typeof ui?.seed_source_type === "string" ? ui.seed_source_type : null; const sourceType = typeof source?.seed_type === "string" ? source.seed_type : null; - const shouldResetHfState = - sourceType === "hf" || uiSourceType === "hf"; + const shouldResetHfState = sourceType === "hf" || uiSourceType === "hf"; const shouldResetLocalState = sourceType === "local" || sourceType === "unstructured" || @@ -144,6 +145,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -165,6 +167,7 @@ function sanitizeSeedForShare(payload: unknown): unknown { ui.seed_drop_columns = []; ui.seed_preview_rows = []; ui.local_file_name = ""; + ui.unstructured_upload_uid = ""; ui.unstructured_file_ids = []; ui.unstructured_file_names = []; ui.unstructured_file_sizes = []; @@ -174,6 +177,43 @@ function sanitizeSeedForShare(payload: unknown): unknown { return root; } +// Delete queued upload directories once a save stops referencing them, so a +// reload before autosave can never leave the saved recipe pointing at +// already-deleted files. Skips any uid the just-saved payload still uses. +function drainQueuedUploadCleanups( + savedPayload: RecipePayloadResult["payload"], +): void { + const pending = useRecipeStudioStore.getState().pendingUploadCleanups; + if (pending.length === 0) { + return; + } + const ui = + savedPayload && typeof savedPayload === "object" + ? (savedPayload as { ui?: Record }).ui + : undefined; + const savedUid = + ui && typeof ui.unstructured_upload_uid === "string" + ? ui.unstructured_upload_uid + : ""; + const ready = pending.filter((uid) => uid !== savedUid); + if (ready.length === 0) { + return; + } + for (const uid of ready) { + void removeUnstructuredBlock(uid) + .then(() => { + useRecipeStudioStore.setState((state) => ({ + pendingUploadCleanups: state.pendingUploadCleanups.filter( + (pendingUid) => pendingUid !== uid, + ), + })); + }) + .catch((error) => { + console.warn("Failed to clean up uploaded documents:", error); + }); + } +} + export function useRecipePersistence({ recipeId, initialRecipeName, @@ -202,8 +242,10 @@ export function useRecipePersistence({ () => buildSignature(normalizedWorkflowName, currentPayload), [currentPayload, normalizedWorkflowName], ); - const isDirty = savedSignature.length > 0 && currentSignature !== savedSignature; - const saveTone: SaveTone = !isDirty && Boolean(lastSavedAt) ? "success" : "error"; + const isDirty = + savedSignature.length > 0 && currentSignature !== savedSignature; + const saveTone: SaveTone = + !isDirty && Boolean(lastSavedAt) ? "success" : "error"; const savedAtLabel = formatSavedLabel(lastSavedAt); useEffect(() => { @@ -214,7 +256,9 @@ export function useRecipePersistence({ setLastSavedAt(initialSavedAt); setCopied(false); - const parsed = importRecipePayload(JSON.stringify(initialPayload)); + const parsed = importRecipePayload(JSON.stringify(initialPayload), { + preserveUnstructuredUploads: true, + }); if (parsed.snapshot) { loadRecipe(parsed.snapshot); } else { @@ -252,6 +296,7 @@ export function useRecipePersistence({ }); setLastSavedAt(result.updatedAt); setSavedSignature(buildSignature(nextName, currentPayload)); + drainQueuedUploadCleanups(currentPayload); } catch (error) { console.error("Save recipe failed:", error); toastError("Save failed", "Could not save recipe."); @@ -270,11 +315,28 @@ export function useRecipePersistence({ return () => window.clearTimeout(timeoutId); }, [isDirty, persistRecipe, saveLoading]); + // Drain queued cleanups even when autosave is skipped: a net-zero edit (add + // then remove an unstructured seed before the 800ms debounce) keeps isDirty + // false, so the autosave effect never drains and the queued uid leaks its + // upload dir. Not-dirty means currentPayload equals the saved recipe, and + // drain skips the uid it still references, so only dirs no saved recipe + // points at are deleted (keeps the save-first invariant). + useEffect(() => { + if (!initialRecipeReady || isDirty || saveLoading) { + return; + } + drainQueuedUploadCleanups(currentPayload); + }, [currentPayload, initialRecipeReady, isDirty, saveLoading]); + const copyRecipe = useCallback(async (): Promise => { setCopied(false); try { - const safePayload = sanitizeSeedForShare(stripApiKeys(payloadResult.payload)); - const ok = await copyTextToClipboard(JSON.stringify(safePayload, null, 2)); + const safePayload = sanitizeSeedForShare( + stripApiKeys(payloadResult.payload), + ); + const ok = await copyTextToClipboard( + JSON.stringify(safePayload, null, 2), + ); if (!ok) { throw new Error("Clipboard not available."); } diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index 8659cad4fb..a1ff72ee9b 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -36,6 +36,7 @@ import { } from "../utils/handles"; import type { RecipeSnapshot } from "../utils/import"; import { getLayoutedElements } from "../utils/layout"; +import { makeUnstructuredUploadUid } from "../utils/config-factories"; import { centerModelInfraNodes, optimizeModelInfraEdgeHandles, @@ -76,6 +77,12 @@ type RecipeStudioState = { nextId: number; nextY: number; fitViewTick: number; + // Upload-uid directories whose owning block dropped them; server-side + // deletion is deferred until a save no longer references them, so a + // reload before autosave cannot leave a saved recipe pointing at + // deleted files. + pendingUploadCleanups: string[]; + queueUploadCleanup: (uid: string) => void; setSheetOpen: (open: boolean) => void; setSheetView: (view: SheetView) => void; setProcessors: (processors: RecipeProcessorConfig[]) => void; @@ -137,6 +144,7 @@ const INITIAL_STATE = { nextId: 3, nextY: 280, fitViewTick: 0, + pendingUploadCleanups: [], } satisfies Pick< RecipeStudioState, | "nodes" @@ -154,6 +162,7 @@ const INITIAL_STATE = { | "nextId" | "nextY" | "fitViewTick" + | "pendingUploadCleanups" >; function buildAddedNodeState( @@ -269,6 +278,20 @@ function isModelSemanticEdge( ); } +// Upload uid of a seed block whose server-side directory becomes orphaned +// when the block drops it. Only uid directories qualify (single owner); +// legacy node-id directories can be shared by other recipes. +function seedUploadCleanupUid(config: NodeConfig | undefined): string | null { + if (!config || config.kind !== "seed") { + return null; + } + const uid = config.unstructured_upload_uid?.trim(); + if (!uid || !config.unstructured_file_ids?.length) { + return null; + } + return uid; +} + export const useRecipeStudioStore = create((set, get) => ({ ...INITIAL_STATE, setSheetOpen: (open) => set({ sheetOpen: open }), @@ -278,6 +301,12 @@ export const useRecipeStudioStore = create((set, get) => ({ setDialogOpen: (open) => set({ dialogOpen: open }), setExecutionLocked: (locked) => set({ executionLocked: locked }), resetRecipe: () => set(INITIAL_STATE), + queueUploadCleanup: (uid) => + set((state) => + state.pendingUploadCleanups.includes(uid) + ? state + : { pendingUploadCleanups: [...state.pendingUploadCleanups, uid] }, + ), selectConfig: (id) => set({ activeConfigId: id, dialogOpen: false }), openConfig: (id) => set({ activeConfigId: id, dialogOpen: true }), setLayoutDirection: (direction) => @@ -383,7 +412,18 @@ export const useRecipeStudioStore = create((set, get) => ({ } return buildAddedNodeState(state, "sampler", type, position, openDialog); }), - addSeedNode: (type, position, openDialog = true) => + addSeedNode: (type, position, openDialog = true) => { + const current = get(); + if (!current.executionLocked) { + // The reset below clears the block's upload uid and file list; queue + // its server-side directory for deletion after the next save. + const uid = seedUploadCleanupUid( + Object.values(current.configs).find((config) => config.kind === "seed"), + ); + if (uid) { + current.queueUploadCleanup(uid); + } + } set((state) => { if (state.executionLocked) { return state; @@ -413,6 +453,8 @@ export const useRecipeStudioStore = create((set, get) => ({ hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + unstructured_upload_uid: + nextSourceType === "unstructured" ? makeUnstructuredUploadUid() : "", unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], @@ -446,7 +488,8 @@ export const useRecipeStudioStore = create((set, get) => ({ activeConfigId: existing.id, dialogOpen: openDialog, }; - }), + }); + }, addLlmNode: (type, position, openDialog = true) => set((state) => { if (state.executionLocked) { @@ -699,6 +742,9 @@ export const useRecipeStudioStore = create((set, get) => ({ dialogOpen: false, sheetView: "root", fitViewTick: state.fitViewTick + 1, + // Queued cleanups belong to the previous recipe; draining them after + // a save of this one could delete files its saved payload still uses. + pendingUploadCleanups: [], })), setAuxNodePosition: (id, position) => set((state) => { @@ -786,6 +832,17 @@ export const useRecipeStudioStore = create((set, get) => ({ set(applyUpdate); }, onNodesChange: (changes) => { + const current = get(); + if (!current.executionLocked) { + for (const change of changes) { + if (change.type === "remove") { + const uid = seedUploadCleanupUid(current.configs[change.id]); + if (uid) { + current.queueUploadCleanup(uid); + } + } + } + } const applyNodesChange = (state: RecipeStudioState) => { if (state.executionLocked) { return state; diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index b8ed13f70b..9231c5f7a3 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -340,6 +340,8 @@ export type SeedConfig = { hf_token?: string; hf_endpoint?: string; local_file_name?: string; + // ui-only: stable per-block id for uploads, since node ids collide across imports + unstructured_upload_uid?: string; unstructured_file_ids?: string[]; unstructured_file_names?: string[]; unstructured_file_sizes?: number[]; diff --git a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts index 76fc2c38ee..d47bac8858 100644 --- a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts +++ b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts @@ -20,6 +20,46 @@ import type { } from "../types"; import { nextName } from "./naming"; +export function makeUnstructuredUploadUid(): string { + if (typeof globalThis.crypto?.randomUUID === "function") { + return globalThis.crypto.randomUUID().replace(/-/g, "").toLowerCase(); + } + if (typeof globalThis.crypto?.getRandomValues === "function") { + const bytes = new Uint8Array(16); + globalThis.crypto.getRandomValues(bytes); + return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join( + "", + ); + } + let uid = ""; + while (uid.length < 32) { + uid += Math.floor(Math.random() * 0x100000000) + .toString(16) + .padStart(8, "0"); + } + return uid.slice(0, 32); +} + +export function resolveUnstructuredUploadBlockId({ + configId, + uploadUid, + generatedUploadUid, + unstructuredFileCount, +}: { + configId: string; + uploadUid: string; + generatedUploadUid: string | null; + unstructuredFileCount: number; +}): string { + if (uploadUid) { + return uploadUid; + } + if (generatedUploadUid) { + return generatedUploadUid; + } + return unstructuredFileCount > 0 ? configId : ""; +} + export function makeSamplerConfig( id: string, samplerType: SamplerType, @@ -368,6 +408,9 @@ export function makeSeedConfig( hf_token: "", hf_endpoint: "https://huggingface.co", local_file_name: "", + ...(seedSourceType === "unstructured" + ? { unstructured_upload_uid: makeUnstructuredUploadUid() } + : {}), unstructured_file_ids: [], unstructured_file_names: [], unstructured_file_sizes: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts index ac881d0373..54df2ffd50 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts @@ -16,11 +16,7 @@ import type { } from "../../types"; import { buildEdges } from "./edges"; import { isRecord, parseJson, readString } from "./helpers"; -import { - parseColumn, - parseModelConfig, - parseModelProvider, -} from "./parsers"; +import { parseColumn, parseModelConfig, parseModelProvider } from "./parsers"; import { parseSeedConfig } from "./parsers/seed-config-parser"; import { buildNodes, parseUi } from "./ui"; import type { ImportResult } from "./types"; @@ -43,6 +39,7 @@ type UiInput = { seed_drop_columns?: unknown; seed_preview_rows?: unknown; local_file_name?: unknown; + unstructured_upload_uid?: unknown; unstructured_file_ids?: unknown; unstructured_file_names?: unknown; unstructured_file_sizes?: unknown; @@ -51,6 +48,10 @@ type UiInput = { advanced_open_by_node?: unknown; }; +type ImportRecipePayloadOptions = { + preserveUnstructuredUploads?: boolean; +}; + type UiMarkdownNoteNode = { name: string; markdown: string; @@ -90,7 +91,7 @@ function parseProcessors(input: unknown): RecipeProcessorConfig[] { ? templateRaw : isRecord(templateRaw) ? JSON.stringify(templateRaw, null, 2) - : "{\n \"text\": \"{{ column_name }}\"\n}"; + : '{\n "text": "{{ column_name }}"\n}'; processors.push({ id: `p${index + 1}`, // biome-ignore lint/style/useNamingConvention: api schema @@ -135,9 +136,7 @@ function parseSeedDropColumns(input: unknown): string[] { return Array.from(values); } -function parseMcpProviders( - input: unknown, -): Map { +function parseMcpProviders(input: unknown): Map { const providers = new Map(); if (!Array.isArray(input)) { return providers; @@ -156,13 +155,12 @@ function parseMcpProviders( const args = Array.isArray(item.args) ? item.args.map((value) => String(value)) : []; - const envPairs = - isRecord(item.env) - ? Object.entries(item.env).map(([key, value]) => ({ - key: String(key), - value: String(value), - })) - : []; + const envPairs = isRecord(item.env) + ? Object.entries(item.env).map(([key, value]) => ({ + key: String(key), + value: String(value), + })) + : []; providers.set(name, { id: `mcp-${index + 1}`, name, @@ -209,7 +207,8 @@ function parseToolConfigs(input: unknown): Map { allow_tools: allowTools, // biome-ignore lint/style/useNamingConvention: api schema max_tool_call_turns: - item.max_tool_call_turns === null || item.max_tool_call_turns === undefined + item.max_tool_call_turns === null || + item.max_tool_call_turns === undefined ? "5" : String(item.max_tool_call_turns), // biome-ignore lint/style/useNamingConvention: api schema @@ -257,7 +256,9 @@ function parseUiMarkdownNoteNodes(input: unknown): UiMarkdownNoteNode[] { return noteNodes; } -function parseUiToolProfileNodes(input: unknown): Map> { +function parseUiToolProfileNodes( + input: unknown, +): Map> { const toolProfiles = new Map>(); if (!Array.isArray(input)) { return toolProfiles; @@ -312,9 +313,15 @@ function parseAdvancedOpenByNode(input: unknown): Record { return out; } -type AdvancedOpenConfig = LlmConfig | SamplerConfig | SeedConfig | ValidatorConfig; +type AdvancedOpenConfig = + | LlmConfig + | SamplerConfig + | SeedConfig + | ValidatorConfig; -function isAdvancedOpenConfig(config: NodeConfig): config is AdvancedOpenConfig { +function isAdvancedOpenConfig( + config: NodeConfig, +): config is AdvancedOpenConfig { return ( config.kind === "llm" || config.kind === "sampler" || @@ -350,7 +357,8 @@ function buildToolProfileConfig( .map((providerName) => mcpProvidersByName.get(providerName)) .flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])), // biome-ignore lint/style/useNamingConvention: ui schema - fetched_tools_by_provider: fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, + fetched_tools_by_provider: + fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, // biome-ignore lint/style/useNamingConvention: api schema allow_tools: [...(canonical.allow_tools ?? [])], // biome-ignore lint/style/useNamingConvention: api schema @@ -360,7 +368,10 @@ function buildToolProfileConfig( }; } -export function importRecipePayload(input: string): ImportResult { +export function importRecipePayload( + input: string, + options: ImportRecipePayloadOptions = {}, +): ImportResult { const parsed = parseJson(input); if (!parsed.data || !isRecord(parsed.data)) { return { @@ -369,9 +380,9 @@ export function importRecipePayload(input: string): ImportResult { }; } - const recipe = (isRecord(parsed.data.recipe) - ? parsed.data.recipe - : parsed.data) as RecipeInput; + const recipe = ( + isRecord(parsed.data.recipe) ? parsed.data.recipe : parsed.data + ) as RecipeInput; const ui = isRecord(parsed.data.ui) ? (parsed.data.ui as UiInput) : null; if (!Array.isArray(recipe.columns)) { @@ -410,21 +421,36 @@ export function importRecipePayload(input: string): ImportResult { .map((row) => ({ ...row })) : undefined; const uiLocalFileName = readString(ui?.local_file_name) ?? undefined; - // Preserve file IDs/names from saved recipes (cleared at share time by sanitizeSeedForShare) - const uiUnstructuredFileIds: string[] = Array.isArray(ui?.unstructured_file_ids) - ? (ui.unstructured_file_ids as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileNames: string[] = Array.isArray(ui?.unstructured_file_names) - ? (ui.unstructured_file_names as string[]).filter((v): v is string => typeof v === "string") - : []; - const uiUnstructuredFileSizes: number[] = Array.isArray(ui?.unstructured_file_sizes) - ? (ui.unstructured_file_sizes as number[]).filter((v): v is number => typeof v === "number") - : []; + const preserveUnstructuredUploads = + options.preserveUnstructuredUploads === true; + const uiUnstructuredUploadUid = preserveUnstructuredUploads + ? (readString(ui?.unstructured_upload_uid) ?? undefined) + : undefined; + const uiUnstructuredFileIds: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_ids) + ? (ui.unstructured_file_ids as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileNames: string[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_names) + ? (ui.unstructured_file_names as string[]).filter( + (v): v is string => typeof v === "string", + ) + : []; + const uiUnstructuredFileSizes: number[] = + preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_sizes) + ? (ui.unstructured_file_sizes as number[]).filter( + (v): v is number => typeof v === "number", + ) + : []; const uiUnstructuredChunkSize = readStringNumber(ui?.unstructured_chunk_size); const uiUnstructuredChunkOverlap = readStringNumber( ui?.unstructured_chunk_overlap, ); - const uiAdvancedOpenByNode = parseAdvancedOpenByNode(ui?.advanced_open_by_node); + const uiAdvancedOpenByNode = parseAdvancedOpenByNode( + ui?.advanced_open_by_node, + ); const uiMarkdownNotes = parseUiMarkdownNoteNodes(ui?.nodes); const uiToolProfilesByName = parseUiToolProfileNodes(ui?.nodes); @@ -459,11 +485,13 @@ export function importRecipePayload(input: string): ImportResult { : payloadSeedDropColumns, seed_preview_rows: uiSeedPreviewRows, local_file_name: uiLocalFileName, + unstructuredUploadUid: uiUnstructuredUploadUid, unstructuredFileIds: uiUnstructuredFileIds, unstructuredFileNames: uiUnstructuredFileNames, unstructuredFileSizes: uiUnstructuredFileSizes, unstructured_chunk_size: uiUnstructuredChunkSize, unstructured_chunk_overlap: uiUnstructuredChunkOverlap, + preserveUnstructuredUploads, }); if (seedConfig) { applyAdvancedOpen(seedConfig, uiAdvancedOpenByNode); @@ -567,12 +595,7 @@ export function importRecipePayload(input: string): ImportResult { const { layouts, auxNodes, edges: uiEdges, layoutDirection } = parseUi(ui); const resolvedLayoutDirection = layoutDirection ?? "LR"; const nodes = buildNodes(configs, layouts); - const edges = buildEdges( - configs, - nameToId, - uiEdges, - resolvedLayoutDirection, - ); + const edges = buildEdges(configs, nameToId, uiEdges, resolvedLayoutDirection); const auxNodePositions = Object.fromEntries( auxNodes.flatMap((item) => { const llmId = nameToId.get(item.llm); @@ -583,10 +606,7 @@ export function importRecipePayload(input: string): ImportResult { }), ); - const maxY = nodes.reduce( - (acc, node) => Math.max(acc, node.position.y), - 0, - ); + const maxY = nodes.reduce((acc, node) => Math.max(acc, node.position.y), 0); return { errors: [], diff --git a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts index 21eadb3195..939205fe6d 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/parsers/seed-config-parser.ts @@ -197,17 +197,26 @@ export function parseSeedConfig( seed_drop_columns?: string[]; seed_preview_rows?: Record[]; local_file_name?: string; + unstructuredUploadUid?: string; unstructuredFileIds?: string[]; unstructuredFileNames?: string[]; unstructuredFileSizes?: number[]; unstructured_chunk_size?: string; unstructured_chunk_overlap?: string; + preserveUnstructuredUploads?: boolean; }, ): SeedConfig | null { if (!seedConfigRaw) { return null; } - const parsed = parseSeedSettings(seedConfigRaw); + const parsed = { ...parseSeedSettings(seedConfigRaw) }; + if ( + parsed.seed_source_type === "unstructured" && + options?.preserveUnstructuredUploads !== true + ) { + parsed.hf_path = ""; + parsed.resolved_paths = []; + } let sourceType: SeedSourceType = "hf"; if (parsed.seed_source_type === "hf") { sourceType = "hf"; @@ -230,6 +239,9 @@ export function parseSeedConfig( ...(options?.local_file_name !== undefined ? { local_file_name: options.local_file_name } : {}), + ...(options?.unstructuredUploadUid + ? { unstructured_upload_uid: options.unstructuredUploadUid } + : {}), ...(options?.unstructuredFileIds !== undefined ? { unstructured_file_ids: options.unstructuredFileIds } : {}), diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts index 34c3c34274..9b2ad5b2b5 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts @@ -440,6 +440,9 @@ export function buildRecipePayload( unstructured_file_names: firstSeed.unstructured_file_names, unstructured_file_sizes: firstSeed.unstructured_file_sizes, }), + ...(firstSeed?.unstructured_upload_uid?.trim() && { + unstructured_upload_uid: firstSeed.unstructured_upload_uid, + }), ...(firstSeed && firstSeed.unstructured_chunk_size !== undefined && { unstructured_chunk_size: firstSeed.unstructured_chunk_size, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts index 902ea796b4..763e68c1fb 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts @@ -71,6 +71,8 @@ export type RecipePayload = { seed_preview_rows?: Record[]; local_file_name?: string; // biome-ignore lint/style/useNamingConvention: api schema + unstructured_upload_uid?: string; + // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_ids?: string[]; // biome-ignore lint/style/useNamingConvention: api schema unstructured_file_names?: string[]; From 1b825213ea2ffe4774404f27795ad287634f6681 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Thu, 9 Jul 2026 00:16:05 -0700 Subject: [PATCH 039/402] Stabilize floating monitor drag (#6984) * Stabilize floating monitor drag * Restore floating monitor exit animation * Harden Windows Studio smoke checks * Keep API menu badge removed * Apply no-build-tools env overrides in-script The runner does not apply step-level env keys containing parentheses, so ProgramFiles(x86) kept its real value and Find-VsBuildTools still detected VS through vswhere. Set the overrides inside each pwsh step instead; child processes inherit them. The resolver step moves to pwsh because bash cannot export a variable named ProgramFiles(x86). * Reset chat UI session without a second browser context macOS runs Chromium with --single-process, where closing the last context tears down the whole browser, so the shutdown re-login died with TargetClosedError on new_page. Clear cookies and swap pages inside the same context instead, opening the replacement page before closing the old one. * Keep the no-build-tools Path filtered across session refreshes install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment rebuild the session Path from the Machine and User registry scopes, so the process-level filter could be undone mid-install and re-expose CMake. Filter those scopes in the Prepare step with normalized dir matching and restore them in cleanup. * Drop stale localStorage auth tokens before re-login Auth tokens live in localStorage, not cookies, and the login guest guard redirects on their mere presence. Remove them during the session reset so the /login navigation is deterministic instead of relying on the tolerated redirect bounce. --- .../studio-windows-inference-smoke.yml | 161 +++++++++---- .../frontend/src/components/app-sidebar.tsx | 17 +- .../src/components/floating-monitor.tsx | 223 ++++++++++-------- studio/frontend/src/features/chat/index.ts | 1 + .../frontend/src/features/settings/index.ts | 1 + studio/frontend/src/i18n/locales/en.ts | 1 - studio/frontend/src/i18n/locales/ja.ts | 1 - studio/frontend/src/i18n/locales/pt-br.ts | 1 - studio/frontend/src/i18n/locales/zh-CN.ts | 1 - tests/studio/playwright_chat_ui.py | 36 ++- 10 files changed, 277 insertions(+), 166 deletions(-) diff --git a/.github/workflows/studio-windows-inference-smoke.yml b/.github/workflows/studio-windows-inference-smoke.yml index 0bc216d65a..dbb0f9ea6f 100644 --- a/.github/workflows/studio-windows-inference-smoke.yml +++ b/.github/workflows/studio-windows-inference-smoke.yml @@ -1334,42 +1334,75 @@ jobs: try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { } } - - name: Hide Visual Studio + CMake (simulate a host with no build tools) + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # A Program Files dir can hold a transient handle (Defender / MSBuild node) - # so Rename-Item intermittently fails with "Access is denied"; retry to ride it out. - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { + [void] $blocked.Add( + [Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\')) + } + } } } - # Rename the Visual Studio install roots (incl. the Installer that holds - # vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss. - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { - Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff') - Write-Host "Hid VS: $d" - } + # Normalized comparison so registry spellings (trailing slash, + # unexpanded %VAR%) still match. + function Test-Blocked([string]$p) { + $n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\') + return $blocked.Contains($n) } - # Surgically rename each cmake executable on PATH (not its parent dir -- - # cmake can share a dir with other shims) so Get-Command cmake fails. - $hidden = @() - foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) { - if ($c.Source -and (Test-Path -LiteralPath $c.Source)) { - Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off') - $hidden += $c.Source - Write-Host "Hid cmake: $($c.Source)" - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not (Test-Blocked $_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + # install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment + # rebuild the session Path from these scopes mid-install, so filter + # them too. Originals are saved for the cleanup step. + foreach ($scope in @('Machine', 'User')) { + $orig = [Environment]::GetEnvironmentVariable('Path', $scope) + if (-not $orig) { continue } + Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline + $kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';' + [Environment]::SetEnvironmentVariable('Path', $kept, $scope) + Write-Host "Filtered $scope Path scope." + } + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH<&1 | Tee-Object -FilePath logs/install.log @@ -1480,19 +1517,19 @@ jobs: [ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; } echo "Inference OK without Visual Studio: $CONTENT" - - name: Restore Visual Studio + CMake + - name: Clean no-build-tools simulation if: always() shell: pwsh run: | - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - $off = "$d.vsoff" - if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } - } - if ($env:HIDDEN_CMAKE) { - foreach ($src in ($env:HIDDEN_CMAKE -split '\|')) { - if ($src -and (Test-Path -LiteralPath "$src.off")) { Rename-Item -LiteralPath "$src.off" -NewName (Split-Path $src -Leaf) } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + foreach ($scope in @('Machine', 'User')) { + $saved = Join-Path $root "orig-path-$scope.txt" + if (Test-Path -LiteralPath $saved) { + [Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope) + Write-Host "Restored $scope Path scope." } } + Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue - name: Stop Studio if: always() @@ -1540,21 +1577,34 @@ jobs: with: python-version: '3.12' - - name: Hide Visual Studio + - name: Prepare no-build-tools simulation shell: pwsh run: | $ErrorActionPreference = 'Stop' - # Retry the rename: a Program Files dir can hold a transient handle that - # makes Rename-Item intermittently fail with "Access is denied". - function Rename-WithRetry($Path, $NewName) { - for ($i = 1; $i -le 6; $i++) { - try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return } - catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 } + $root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools' + $pf = Join-Path $root 'ProgramFiles' + $pfx86 = Join-Path $root 'ProgramFilesx86' + New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null + + $blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase) + foreach ($tool in @('cmake', 'cl.exe')) { + foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) { + if ($cmd.Source) { + $dir = Split-Path -Parent $cmd.Source + if ($dir) { [void] $blocked.Add($dir) } + } } } - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" } - } + + $pathParts = $env:Path -split [IO.Path]::PathSeparator | + Where-Object { $_ -and -not $blocked.Contains($_) } + $noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator + + "NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8 + "NO_BUILD_TOOLS_PATH< /tmp/resolve.json || { - echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; } - cat /tmp/resolve.json - echo "Prebuilt resolver ran with no Visual Studio present." + if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 } + python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json + if ($LASTEXITCODE -ne 0) { + Write-Host "::error::resolver exited non-zero" + if (Test-Path resolve.json) { Get-Content resolve.json } + exit 1 + } + Get-Content resolve.json + Write-Host "Prebuilt resolver ran with no Visual Studio present." - - name: Restore Visual Studio + - name: Clean no-build-tools simulation if: always() shell: pwsh run: | - foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) { - $off = "$d.vsoff" - if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" } - } + Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue # ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ── pester: diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 823e420869..1a74b38524 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -81,7 +81,6 @@ import { TestTube01Icon, ZapIcon, } from "@hugeicons/core-free-icons"; -import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage"; import { Tooltip, TooltipContent, @@ -97,6 +96,7 @@ import { createChatProject, deleteChatProject, deleteChatItem, + listStoredChatThreads, moveChatItemToProject, renameChatItem, renameChatProject, @@ -582,7 +582,14 @@ export function AppSidebar() { useEffect(() => { if (!pendingRename) return; const match = allChatItems.find((i) => i.id === pendingRename.id); - if (match && match.title === pendingRename.title) setPendingRename(null); + if (!match || match.title !== pendingRename.title) return; + queueMicrotask(() => { + setPendingRename((current) => + current?.id === pendingRename.id && current.title === pendingRename.title + ? null + : current, + ); + }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); const [projectNameDraft, setProjectNameDraft] = useState(""); @@ -680,12 +687,6 @@ export function AppSidebar() { useState(null); const [deleteProjectFiles, setDeleteProjectFiles] = useState(false); - useEffect(() => { - if (confirmingDelete?.kind !== "project") { - setDeleteProjectFiles(false); - } - }, [confirmingDelete]); - async function commitDelete() { const target = confirmingDelete; if (!target) return; diff --git a/studio/frontend/src/components/floating-monitor.tsx b/studio/frontend/src/components/floating-monitor.tsx index bce4bf2831..0a51875de9 100644 --- a/studio/frontend/src/components/floating-monitor.tsx +++ b/studio/frontend/src/components/floating-monitor.tsx @@ -3,27 +3,35 @@ import { Button } from "@/components/ui/button"; import { Progress } from "@/components/ui/progress"; -import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store"; +import { useMonitorOverlayStore } from "@/features/settings"; import { useSystemInfo } from "@/hooks/use-system"; import { useT } from "@/i18n"; import { cn } from "@/lib/utils"; import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react"; -import { motion } from "motion/react"; -import { useRef } from "react"; +import { AnimatePresence, motion, useDragControls } from "motion/react"; +import { type PointerEvent, useMemo, useState } from "react"; function clampPercent(value: number): number { return Math.max(0, Math.min(100, value)); } function usageIndicatorClass(percent: number): string { - if (percent >= 90) return "bg-destructive"; - if (percent >= 70) return "bg-amber-500"; + if (percent >= 90) { + return "bg-destructive"; + } + if (percent >= 70) { + return "bg-amber-500"; + } return "bg-primary"; } function usageTextClass(percent: number): string { - if (percent >= 90) return "text-destructive"; - if (percent >= 70) return "text-amber-600 dark:text-amber-400"; + if (percent >= 90) { + return "text-destructive"; + } + if (percent >= 70) { + return "text-amber-600 dark:text-amber-400"; + } return "text-primary"; } @@ -39,9 +47,18 @@ export function FloatingMonitor() { const { isOpen, setIsOpen } = useMonitorOverlayStore(); const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 }); - const constraintsRef = useRef(null); + const [constraintsElement, setConstraintsElement] = + useState(null); + const constraintsRef = useMemo( + () => ({ current: constraintsElement }), + [constraintsElement], + ); + const dragControls = useDragControls(); - if (!isOpen) return null; + function startDrag(event: PointerEvent) { + event.preventDefault(); + dragControls.start(event); + } const ramTotal = systemInfo.memory?.total_gb ?? 0; const ramAvailable = systemInfo.memory?.available_gb ?? 0; @@ -64,99 +81,109 @@ export function FloatingMonitor() { const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0; return ( -
- -
-
- - - {t("settings.resources.liveMonitor.title")} - -
-
-
- -
- - -
-
- - + {isOpen && ( +
-
-
- {t("settings.resources.liveMonitor.ram")} - - {Math.round(ramPercent)}% - -
-
- {formatGiB(ramUsed)} / {formatGiB(ramTotal)} -
- -
- - {hasGpu && ( -
-
- - {t("settings.resources.liveMonitor.vram")}{" "} - {devices.length > 1 - ? `(${devices.length} GPUs)` - : `(${devices[0].name ?? "GPU"})`} + +
+
+ + + {t("settings.resources.liveMonitor.title")} - +
+
- {Math.round(vramPercent)}% - + +
+ +
-
- {formatGiB(vramUsed)} / {formatGiB(vramTotal)} -
-
- )} - - -
+ + +
+
+ {t("settings.resources.liveMonitor.ram")} + + {Math.round(ramPercent)}% + +
+
+ {formatGiB(ramUsed)} / {formatGiB(ramTotal)} +
+ +
+ + {hasGpu && ( +
+
+ + {t("settings.resources.liveMonitor.vram")}{" "} + {devices.length > 1 + ? `(${devices.length} GPUs)` + : `(${devices[0].name ?? "GPU"})`} + + + {Math.round(vramPercent)}% + +
+
+ {formatGiB(vramUsed)} / {formatGiB(vramTotal)} +
+ +
+ )} + + +
+ )} + ); } diff --git a/studio/frontend/src/features/chat/index.ts b/studio/frontend/src/features/chat/index.ts index 7cd9611c71..d070ed15de 100644 --- a/studio/frontend/src/features/chat/index.ts +++ b/studio/frontend/src/features/chat/index.ts @@ -36,6 +36,7 @@ export { ChatSearchDialog } from "./components/chat-search-dialog"; export { setTrainingCompareHandoff } from "./lib/training-compare-handoff"; export type { ProjectRecord } from "./types"; export { clearAllChats, countAllChats } from "./utils/clear-all-chats"; +export { listStoredChatThreads } from "./utils/chat-history-storage"; export { ArtifactCard } from "./artifacts/artifact-card"; export { useChatArtifactsStore, diff --git a/studio/frontend/src/features/settings/index.ts b/studio/frontend/src/features/settings/index.ts index 364ca1611f..3fefd8c63a 100644 --- a/studio/frontend/src/features/settings/index.ts +++ b/studio/frontend/src/features/settings/index.ts @@ -7,6 +7,7 @@ export { savePersonalization, } from "./api/personalization"; export { setTheme, useTheme } from "./stores/theme-store"; +export { useMonitorOverlayStore } from "./stores/monitor-overlay-store"; export type { Personalization, PersonalizationAppearance, diff --git a/studio/frontend/src/i18n/locales/en.ts b/studio/frontend/src/i18n/locales/en.ts index a9b5d839b3..e0cb8030ae 100644 --- a/studio/frontend/src/i18n/locales/en.ts +++ b/studio/frontend/src/i18n/locales/en.ts @@ -409,7 +409,6 @@ export const en = { description: "Access Unsloth via the OpenAI-compatible API.", readDocs: "Read the API docs", noAccess: "No API access yet.", - newBadge: "New", accessTokens: "Access tokens", loadError: "Couldn't load API access.", createError: "Couldn't create access token.", diff --git a/studio/frontend/src/i18n/locales/ja.ts b/studio/frontend/src/i18n/locales/ja.ts index f2020f3cfb..08b8d4f4d3 100644 --- a/studio/frontend/src/i18n/locales/ja.ts +++ b/studio/frontend/src/i18n/locales/ja.ts @@ -296,7 +296,6 @@ export const ja = { description: "OpenAI互換 API を介して Unsloth にアクセスします。", readDocs: "API ドキュメントを読む", noAccess: "まだ API アクセス権がありません。", - newBadge: "新規", accessTokens: "アクセストークン", loadError: "API アクセス権を読み込めませんでした。", createError: "アクセストークンを作成できませんでした。", diff --git a/studio/frontend/src/i18n/locales/pt-br.ts b/studio/frontend/src/i18n/locales/pt-br.ts index c261e4ed0c..494a98ec32 100644 --- a/studio/frontend/src/i18n/locales/pt-br.ts +++ b/studio/frontend/src/i18n/locales/pt-br.ts @@ -363,7 +363,6 @@ export const ptBR = { "Acesse o Unsloth por meio da API compatível com OpenAI.", readDocs: "Leia a documentação da API", noAccess: "Nenhum acesso à API ainda.", - newBadge: "Novo", accessTokens: "Tokens de acesso", loadError: "Não foi possível carregar o acesso à API.", createError: "Não foi possível criar o token de acesso.", diff --git a/studio/frontend/src/i18n/locales/zh-CN.ts b/studio/frontend/src/i18n/locales/zh-CN.ts index f6dc265fc7..dda8e017ab 100644 --- a/studio/frontend/src/i18n/locales/zh-CN.ts +++ b/studio/frontend/src/i18n/locales/zh-CN.ts @@ -267,7 +267,6 @@ export const zhCN = { description: "通过兼容 OpenAI 的 API 以编程方式访问 Unsloth。", readDocs: "阅读 API 文档", noAccess: "还没有 API 访问权限。", - newBadge: "新", accessTokens: "访问 token", loadError: "无法加载 API 访问权限。", createError: "无法创建访问 token。", diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 71297f9043..0698d8e0d2 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -1264,11 +1264,37 @@ with sync_playwright() as p: # placeholder, and /api/health goes unreachable shortly after. # ───────────────────────────────────────────────────── step("Shutdown via account menu") - # Re-login with NEW2 for a valid /api/shutdown token (CLI rotation - # invalidated the old one). The stale token can make the SPA auth guard - # abort this goto with ERR_ABORTED, or redirect to the same /login URL - # ("interrupted by another navigation"); resolve on domcontentloaded and - # tolerate either -- the pw-field wait below confirms we are on /login. + # Start fresh after the CLI rotation invalidates this browser session. + # Stay in the SAME context: macOS Chromium runs --single-process, where + # closing the last context kills the browser and a second context cannot + # be created. Open the new page before closing the old one; the context + # init script covers the new page. + try: + ctx.clear_cookies() + except Exception as exc: + info(f"WARN clearing stale session cookies failed: {exc!r}") + # Auth tokens live in localStorage, and /login's guest guard redirects on + # their mere presence, so drop them before navigating. + try: + page.evaluate( + "['unsloth_auth_token', 'unsloth_auth_refresh_token']" + ".forEach((key) => localStorage.removeItem(key))" + ) + except Exception as exc: + info(f"WARN clearing stale auth tokens failed: {exc!r}") + _fresh_page = ctx.new_page() + _fresh_page.set_default_timeout(60_000) + _fresh_page.on("pageerror", lambda e: page_errors.append(str(e))) + _fresh_page.on("console", _on_console) + try: + page.close() + except Exception: + pass + page = _fresh_page + + # Re-login with NEW2 for a valid /api/shutdown token. Route changes can + # still abort or interrupt this navigation, so the field wait below is the + # final confirmation that we reached /login. _tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation") try: page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000) From 8205d4c0819088a3c864fdd505ce1c1e6d72852d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 01:46:14 -0700 Subject: [PATCH 040/402] Retry the Studio UI shutdown re-login on transient goto timeout (#7027) * Retry the Studio UI shutdown re-login on transient goto timeout The Chat UI Playwright smoke intermittently failed at the pre-shutdown re-login: page.goto('/login') can hit a 60s TimeoutError on a slow runner even while the server is healthy, and the surrounding except only tolerated ERR_ABORTED / interrupted-navigation, so a plain timeout hard-failed the job. Wrap the re-login goto/wait/fill/submit in the same 3-attempt retry the change-password step already uses (recover_or_replace_page between tries, per-attempt fail screenshots, wait_for_health pre-gate). The composer wait stays outside the loop so a retry never re-navigates after login has set tokens (which would redirect to /chat via the guest guard); it remains the authoritative confirmation, so a genuinely broken login still fails. * Catch transient login-request failures and preserve error listeners on recovery Wait on the /api/auth/login POST inside the retry (via click_and_wait_for_response) so a transient 4xx/5xx is retried in-loop instead of surfacing only at the out-of-loop composer wait, matching the change-password step. When recover_or_replace_page swaps in a fresh page, re-attach the pageerror/console listeners so error tracking survives the replacement. --- tests/studio/playwright_chat_ui.py | 96 ++++++++++++++++++++++++++---- 1 file changed, 86 insertions(+), 10 deletions(-) diff --git a/tests/studio/playwright_chat_ui.py b/tests/studio/playwright_chat_ui.py index 0698d8e0d2..6a88b98c19 100644 --- a/tests/studio/playwright_chat_ui.py +++ b/tests/studio/playwright_chat_ui.py @@ -1296,16 +1296,92 @@ with sync_playwright() as p: # still abort or interrupt this navigation, so the field wait below is the # final confirmation that we reached /login. _tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation") - try: - page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000) - except Exception as exc: - if not any(t in str(exc) for t in _tolerated_nav): - raise - info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login") - pw_field = page.locator("#password") - pw_field.wait_for(state = "visible", timeout = 60_000) - pw_field.fill(NEW2) - page.locator('button[type="submit"]').click() + # A slow CI runner can make this re-login navigation time out even with the + # server healthy, so retry the whole goto/wait/fill/submit sequence (mirrors + # the change-password retry above). wait_for_health is a diagnostic pre-gate. + wait_for_health(BASE, timeout = 30.0, info = info) + relogin_err: Exception | None = None + for _relogin_attempt in range(3): + try: + try: + page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000) + except Exception as exc: + if not any(t in str(exc) for t in _tolerated_nav): + raise + info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login") + pw_field = page.locator("#password") + pw_field.wait_for(state = "visible", timeout = 60_000) + pw_field.fill(NEW2) + # Wait on the login POST so a transient 4xx/5xx is caught and retried + # here, not swallowed until the out-of-loop composer wait. + status, _ = click_and_wait_for_response( + page, + url_substr = "/api/auth/login", + method = "POST", + do_click = lambda: page.locator('button[type="submit"]').click(), + timeout_ms = 30_000, + info = lambda m: print(f"[ui] {m}", flush = True), + ) + if status is not None and status >= 400: + raise AssertionError( + f"login POST returned {status}; see console_errors={console_errors[:1]!r}" + ) + relogin_err = None + break + except Exception as e: + relogin_err = e + try: + cur_url = page.url + except Exception: + cur_url = "" + print( + f"[ui] re-login attempt {_relogin_attempt + 1} failed: " + f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; " + f"page_errors={len(page_errors)} console_errors={len(console_errors)}", + flush = True, + ) + if console_errors: + print( + f"[ui] first console.error: {console_errors[0][:200]!r}", + flush = True, + ) + if page_errors: + print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True) + try: + shoot(f"18-relogin-attempt-{_relogin_attempt + 1}-fail") + except Exception: + pass + if _relogin_attempt < 2: + # ERR_NO_BUFFER_SPACE needs the OS to recover socket + # buffers; back off 5s then 15s before retrying. + if "ERR_NO_BUFFER_SPACE" in str(e): + backoff_s = 5 if _relogin_attempt == 0 else 15 + print( + f"[ui] ENOBUFS detected; sleeping {backoff_s}s " + f"before retry to let OS recover socket buffers...", + flush = True, + ) + time.sleep(backoff_s) + # Replace the page if it died; otherwise next iteration's + # page.goto() handles the reload. + old_page = page + page = recover_or_replace_page( + page, + ctx, + default_timeout_ms = 60_000, + info = lambda m: print(f"[ui] recovery: {m}", flush = True), + ) + # A freshly created replacement page loses the pageerror/console + # listeners; re-attach so error tracking survives recovery. + if page is not old_page: + page.on("pageerror", lambda e: page_errors.append(str(e))) + page.on("console", _on_console) + if relogin_err is not None: + raise relogin_err + # Composer mount confirms the rotated session is authenticated. Kept OUTSIDE the + # retry: the loop breaks right after submit, so we never re-goto /login once login + # has set tokens -- that would hit the guest guard, redirect to /chat, and make a + # merely-slow composer look like a broken login. composer = page.locator('textarea[aria-label="Message input"]') composer.wait_for(state = "visible", timeout = 60_000) shoot("18-relogin-with-NEW2") From 5e43c623b98affc23efbf9dbe71061de7c1706a2 Mon Sep 17 00:00:00 2001 From: Etherl <61019402+Etherll@users.noreply.github.com> Date: Thu, 9 Jul 2026 11:46:22 +0300 Subject: [PATCH 041/402] Fix FastSentenceTransformer Qwen embedding preprocessing (#6939) * Fix FastSentenceTransformer Qwen embedding preprocessing * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Document Transformer.load embedding modality fix for #6881 * Harden #6881 fix and add forwards/backwards-compatible regression tests * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fall back to Transformer constructor on legacy sentence-transformers without Hub-capable load * Mirror legacy sentence-transformers fallback in embedding-parity tripwire test * Tighten #6881 comments and docstrings * Skip embedding-parity test on CPU-only runners since FastSentenceTransformer requires CUDA * Honor the transformer module's saved subfolder when loading modules.json records a path for the Transformer module (root for decoder embedders like Qwen3-Embedding, 0_Transformer for the classic layout). Pooling/Normalize already load from their saved path; thread the same path into Transformer.load as subfolder so config and tokenizer resolve like stock ST. stays a no-op, so single-module models are unchanged. * Make embedding-parity test bf16-aware fp16 overflows to NaN on bf16-native embedders such as EmbeddingGemma (Gemma3), producing a false parity failure. Prefer bf16 when the GPU supports it so the tripwire can guard the full documented embedding matrix (Qwen3-Embedding, EmbeddingGemma, BGE-M3, all-MiniLM, GTE-ModernBERT), not just fp16-safe models. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: danielhanchen --- ...t_sentence_transformer_embedding_parity.py | 122 ++++++++++++++++++ ...st_sentence_transformers_pinned_symbols.py | 38 ++++++ unsloth/models/sentence_transformer.py | 65 +++++++++- 3 files changed, 222 insertions(+), 3 deletions(-) create mode 100644 tests/python/test_fast_sentence_transformer_embedding_parity.py diff --git a/tests/python/test_fast_sentence_transformer_embedding_parity.py b/tests/python/test_fast_sentence_transformer_embedding_parity.py new file mode 100644 index 0000000000..252d4486a5 --- /dev/null +++ b/tests/python/test_fast_sentence_transformer_embedding_parity.py @@ -0,0 +1,122 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. +"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text +like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a +"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building +via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings; +`_create_transformer_module` uses `Transformer.load(...)` instead. + +Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST +is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity, +opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected). +""" + +from __future__ import annotations + +import inspect +import os + +import pytest + + +def test_transformer_load_signature_supports_unsloth_kwargs(): + """Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs + the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back + to Transformer(...) there, so mirror that gate and skip.""" + models = pytest.importorskip("sentence_transformers.models") + load = getattr(models.Transformer, "load", None) + assert callable(load), ( + "sentence_transformers Transformer.load is missing; the #6881 fix in " + "unsloth.models.sentence_transformer._create_transformer_module depends on it." + ) + params = inspect.signature(load).parameters + accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values()) + # Mirror _create_transformer_module's hub_capable gate. + hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision")) + if not hub_capable: + pytest.skip( + "legacy Transformer.load(input_path); production path falls back to Transformer(...)" + ) + unsupported = [ + k + for k in ("token", "cache_folder", "revision", "trust_remote_code") + if not (accepts_var_kw or k in params) + ] + assert not unsupported, ( + f"installed sentence_transformers Transformer.load no longer accepts {unsupported} " + f"and has no **kwargs; update _create_transformer_module (#6881) before it silently " + f"falls back to Transformer(...)." + ) + + +def _probe_texts(): + return [ + "roasted chickpeas in 20 kg bags", + "The capital of France is Paris.", + "A fast brown fox jumps over the lazy dog.", + "recette de tarte aux pommes traditionnelle", + ] + + +def test_fast_sentence_transformer_matches_stock_st(): + """End-to-end: FastSentenceTransformer embeddings and tokenization must match a + stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and + GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners.""" + model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL") + if not model_id: + pytest.skip( + "set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model " + "(HF id or local path) to run the #6881 parity test" + ) + + torch = pytest.importorskip("torch") + if not torch.cuda.is_available(): + pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner") + np = pytest.importorskip("numpy") + pytest.importorskip("sentence_transformers") + from sentence_transformers import SentenceTransformer + + device = "cuda" + # Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native + # embedders such as EmbeddingGemma (Gemma3), which would mask real parity. + dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 + texts = _probe_texts() + max_seq_length = 256 + + # Control FIRST, before importing unsloth, so its global import patches never + # touch the stock reference (mirrors the issue's "restart runtime" repro). + ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype}) + ctrl.max_seq_length = max_seq_length + ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist() + ctrl_emb = np.asarray( + ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32 + ) + + import unsloth # noqa: F401 + from unsloth import FastSentenceTransformer + + fast = FastSentenceTransformer.from_pretrained( + model_id, + max_seq_length = max_seq_length, + dtype = dtype, + load_in_4bit = False, + load_in_16bit = True, + ) + fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist() + fast_emb = np.asarray( + fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32 + ) + + # Identical tokenization = no chat-template wrapping slipped in (the #6881 defect). + assert fast_ids == ctrl_ids, ( + f"tokenization diverged (chat-template wrapping regressed?):\n" + f" stock: {ctrl_ids}\n fast: {fast_ids}" + ) + + cos = (ctrl_emb * fast_emb).sum(1) / ( + np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1) + ) + assert float(cos.min()) > 0.99, ( + f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 " + f"(per-text {[round(float(c), 5) for c in cos]})" + ) diff --git a/tests/version_compat/test_sentence_transformers_pinned_symbols.py b/tests/version_compat/test_sentence_transformers_pinned_symbols.py index c0c35b9d5d..d7f9a54811 100644 --- a/tests/version_compat/test_sentence_transformers_pinned_symbols.py +++ b/tests/version_compat/test_sentence_transformers_pinned_symbols.py @@ -19,6 +19,8 @@ ST_TAGS = [ "v5.2.3", "v5.3.0", "v5.4.1", + "v5.5.1", + "v5.6.0", "master", ] @@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str): ) +# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881). +@pytest.mark.parametrize("tag", ST_TAGS) +def test_st_transformer_load_accepts_unsloth_kwargs(tag: str): + """unsloth builds saved ST models via Transformer.load(...) so the saved + modality_config is honored (#6881). If .load stops accepting the hub kwargs it + passes (and has no **kwargs), update the fix before it silently regresses. Not + locating .load is a SKIP (may be inherited); the live test guards the install.""" + candidates = [ + "sentence_transformers/models/Transformer.py", + "sentence_transformers/models/transformer.py", + "sentence_transformers/base/modules/transformer.py", + "sentence_transformers/base/modules/module.py", + ] + for p in candidates: + src = fetch_text("UKPLab/sentence-transformers", tag, p) + if src is None or not has_def(src, "load", "func"): + continue + m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S) + if m is None: + continue + sig = m.group(1) + accepts_var_kw = "**" in sig + missing = [ + kw + for kw in ("token", "cache_folder", "revision", "trust_remote_code") + if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig)) + ] + assert not missing, ( + f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no " + f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module " + f"(#6881) before it silently falls back to Transformer(...)." + ) + return + pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)") + + # sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls. @pytest.mark.parametrize("tag", ST_TAGS) def test_st_util_helpers(tag: str): diff --git a/unsloth/models/sentence_transformer.py b/unsloth/models/sentence_transformer.py index c1172faa94..4a1a555bcf 100644 --- a/unsloth/models/sentence_transformer.py +++ b/unsloth/models/sentence_transformer.py @@ -990,7 +990,17 @@ class FastSentenceTransformer(FastModel): return None @staticmethod - def _create_transformer_module(model_name, model, tokenizer, max_seq_length, trust_remote_code): + def _create_transformer_module( + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + token = None, + cache_dir = None, + revision = None, + module_subfolder = "", + ): """Helper to create and configure a Transformer module.""" from sentence_transformers.models import Transformer @@ -1077,7 +1087,45 @@ class FastSentenceTransformer(FastModel): elif "tokenizer_args" in transformer_init_params: transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy() - transformer_module = Transformer(model_name, **transformer_kwargs) + # Build via Transformer.load so the saved modality_config is honored: plain + # Transformer(...) makes ST 5.x infer a "message" modality for chat-template + # models (e.g. Qwen3-Embedding), chat-wrapping inputs and degrading embeddings + # (#6881). Only use .load when it resolves a Hub id (accepts the kwargs or + # **kwargs); legacy ST 3.x/4.x load(input_path) is local-only with no modality + # bug, so fall back to the constructor. + transformer_module = None + transformer_load = getattr(Transformer, "load", None) + has_modules_json = ( + FastSentenceTransformer._module_path( + model_name, token, cache_dir = cache_dir, revision = revision + ) + is not None + ) + if callable(transformer_load) and has_modules_json: + load_params = inspect.signature(transformer_load).parameters + accepts_var_kw = any( + p.kind is inspect.Parameter.VAR_KEYWORD for p in load_params.values() + ) + hub_capable = accepts_var_kw or any( + key in load_params for key in ("token", "cache_folder", "revision") + ) + if hub_capable: + load_kwargs = { + "token": token, + "cache_folder": cache_dir, + "revision": revision, + "trust_remote_code": trust_remote_code, + **transformer_kwargs, + } + # Resolve config/tokenizer from the module's saved subfolder + # (modules.json "path"), like stock ST; "" (root) is a no-op. + if module_subfolder: + load_kwargs["subfolder"] = module_subfolder + if not accepts_var_kw: + load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params} + transformer_module = Transformer.load(model_name, **load_kwargs) + if transformer_module is None: + transformer_module = Transformer(model_name, **transformer_kwargs) finally: # Restore original Auto* loading immediately AutoModel.from_pretrained = original_model_from_pretrained @@ -1191,6 +1239,10 @@ class FastSentenceTransformer(FastModel): tokenizer, max_seq_length, trust_remote_code, + token, + cache_dir, + revision, + module_subfolder = module_config.get("path") or "", ) modules[name] = transformer_module else: @@ -1226,7 +1278,14 @@ class FastSentenceTransformer(FastModel): ) transformer_module = FastSentenceTransformer._create_transformer_module( - model_name, model, tokenizer, max_seq_length, trust_remote_code + model_name, + model, + tokenizer, + max_seq_length, + trust_remote_code, + token, + cache_dir, + revision, ) modules["0"] = transformer_module From 6d674e5cc9aef396ce8aae45306b2b42beb76244 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:08:39 -0700 Subject: [PATCH 042/402] unsloth start: warn before running an agent's remote installer (#7024) When a coding agent is missing, `unsloth start ` offers to run the vendor's own installer (curl | bash, irm | iex, or npm) after an interactive confirm. Those installers execute with the user's privileges and there is no signature or hash check on the fetched content, so a blind "yes" is a supply-chain risk if the delivery path is compromised. Keep the auto-install convenience but make consent informed: before the prompt, name the exact remote source the installer fetches (or the command it runs for a package installer) and state that nothing verifies a signature or hash. Behavior is otherwise unchanged: non-interactive stdin still never executes anything, and the confirm still defaults to no. --- unsloth_cli/commands/start.py | 19 ++++++++++++++++++- unsloth_cli/tests/test_start.py | 26 ++++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 1 deletion(-) diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index 764f5c7963..a8665b9be3 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -989,6 +989,12 @@ def _refresh_windows_path() -> None: os.environ["PATH"] = os.pathsep.join(entries) +def _install_source(install_hint: str) -> Optional[str]: + """The first http(s) URL an install hint fetches, or None (e.g. an npm install).""" + match = re.search(r"https?://[^\s'\")]+", install_hint) + return match.group(0) if match else None + + def _install_agent(name: str, install_hint: str) -> Optional[str]: # Missing agent under --launch: offer to run its documented install command, then # re-resolve it on PATH. Consent-based (we never auto-run a remote install script @@ -997,7 +1003,18 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]: if not sys.stdin.isatty(): return None typer.echo(f"`{name}` is not installed.") - if not typer.confirm(f"Install it now with `{install_hint}`?", default = False): + # Make the supply-chain risk explicit before the prompt: these are the vendors' + # own installers (curl | bash, irm | iex, npm), run with the user's privileges, + # and nothing checks a signature or hash on the fetched content. Naming the source + # turns a blind "yes" into informed consent. + source = _install_source(install_hint) + warning = ( + f"This will download and RUN a script from {source} with your privileges" + if source + else f"This will RUN `{install_hint}` with your privileges" + ) + typer.secho(f"{warning}; there is no signature or hash check.", fg = "yellow", err = True) + if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False): return None # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm) # on Windows, /bin/sh (curl | bash, or npm) everywhere else. diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 18cb40f18d..87a295532e 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -128,6 +128,32 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch): assert ran == [["powershell", "-NoProfile", "-Command", install_hint]] +def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys): + # Before the confirm, a remote installer must name the URL it fetches so the + # user consents to a specific source rather than blindly accepting. + monkeypatch.setattr(start.os, "name", "nt") + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) # decline: nothing runs + hint = "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup" + assert start._install_agent("hermes", hint) is None + err = capsys.readouterr().err + assert "https://hermes-agent.nousresearch.com/install.ps1" in err + assert "download and RUN" in err + assert "signature or hash" in err + + +def test_install_agent_warns_for_package_installer(monkeypatch, capsys): + # An npm-style installer has no URL to fetch, but still runs with the user's + # privileges, so the warning names the command instead. + monkeypatch.setattr(start.os, "name", "posix") + monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True)) + monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) + assert start._install_agent("codex", "npm install -g @openai/codex") is None + err = capsys.readouterr().err + assert "npm install -g @openai/codex" in err + assert "with your privileges" in err + + def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch): monkeypatch.setattr(start.os, "name", "nt") From 0d4bd50768ca1d55009b51dfa097d7c71e819f4b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:26:24 -0700 Subject: [PATCH 043/402] Restore process-global torch.compile config on torch 2.12 so gradient checkpointing backward honors it (#7019) * Mirror dynamo/inductor config sets into defaults so torch 2.12 worker threads honor them torch 2.12 stores config user overrides in ContextVars, so direct assignments like torch._dynamo.config.recompile_limit = 1024 no longer reach the autograd engine worker threads. Gradient checkpointing recomputes fullgraph-compiled gpt-oss kernels inside backward on those threads, which then read the default recompile limit of 8 and raise FailOnRecompileLimitHit at step 0 of GRPO/SFT. Mirror direct config assignments into the process-global entry defaults on torch >= 2.12, restoring the torch <= 2.11 cross-thread semantics while leaving the context-scoped config.patch API untouched. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Keep config.patch thread-local when mirroring dynamo/inductor sets config.patch(...) also assigns through ConfigModule.__setattr__, so the default-mirror was leaking its scoped, thread-local writes into the process-global entry default. Track patch enter/exit with a per-thread depth counter (wrapping ConfigModule.patch) and skip mirroring while inside a patch, so only genuine direct assignments restore the torch 2.11 cross-thread semantics and config.patch stays context-local. * Also keep config.load_config thread-local when mirroring config sets load_config restores a saved dynamo/inductor config by calling setattr per key, which the default-mirror would otherwise leak process-wide just like config.patch did. Wrap load_config with the same per-thread depth counter (renamed to _scoped_depth) so both scoped writers skip the mirror and stay context-local, while genuine direct assignments still restore the torch 2.11 cross-thread default. * Drop the pre-existing override replay from the config thread fix The replay was redundant: this runs from _gpu_init before unsloth sets any dynamo/inductor config, so the __setattr__ wrapper already mirrors every later assignment (recompile_limit included). It could also read a value that belonged to a config.patch context still active at import time and write that thread-local override into the global default. Removing it keeps the cross-thread fix and drops the now-unused _inductor.config import. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/_gpu_init.py | 6 ++ unsloth/import_fixes.py | 132 +++++++++++++++++++++++++++++++++++++++- 2 files changed, 136 insertions(+), 2 deletions(-) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index e39b44488a..e6178e60f3 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -173,6 +173,7 @@ from .import_fixes import ( fix_vllm_guided_decoding_params, fix_vllm_pdl_blackwell, fix_triton_compiled_kernel_missing_attrs, + fix_dynamo_config_thread_visibility, patch_trunc_normal_precision_issue, ignore_logger_messages, patch_ipykernel_hf_xet, @@ -203,6 +204,10 @@ fix_vllm_guided_decoding_params() fix_trl_vllm_ascend() fix_vllm_pdl_blackwell() fix_triton_compiled_kernel_missing_attrs() +# Must run before unsloth_zoo's patch_torch_compile and the gpt-oss temporary +# patches raise the dynamo recompile limits, so those settings reach the +# autograd worker threads on torch >= 2.12. +fix_dynamo_config_thread_visibility() patch_trunc_normal_precision_issue() ignore_logger_messages() patch_ipykernel_hf_xet() @@ -233,6 +238,7 @@ del fix_vllm_guided_decoding_params del fix_trl_vllm_ascend del fix_vllm_pdl_blackwell del fix_triton_compiled_kernel_missing_attrs +del fix_dynamo_config_thread_visibility del patch_trunc_normal_precision_issue del ignore_logger_messages del patch_ipykernel_hf_xet diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index e5a5d01c2f..c5300ed4d0 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -1064,6 +1064,135 @@ def fix_triton_compiled_kernel_missing_attrs(): ) +def fix_dynamo_config_thread_visibility(): + """torch 2.12 made torch._dynamo/_inductor config overrides thread-local + (ContextVars), so `config.recompile_limit = 1024` set on the main thread is + invisible to the autograd worker threads that run backward. Gradient + checkpointing recompiles fullgraph gpt-oss kernels there against the default + limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config + assignments into the process-global entry default (torch <= 2.11 semantics). + config.patch(...) and config.load_config(...) also assign via __setattr__ but + are thread-local by design, so skip mirroring while inside one (tracked per + thread). No-op below torch 2.12 and on any torch without this internal layout. + """ + try: + import torch + + if Version(torch.__version__) < Version("2.12.0"): + return + import torch._dynamo.config as _dynamo_config + from torch.utils._config_module import ConfigModule + from contextvars import ContextVar + except Exception: + return + + try: + probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None) + if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar): + # Overrides are not context-local on this torch; nothing to fix. + return + original_setattr = ConfigModule.__setattr__ + if getattr(original_setattr, "__unsloth_patched__", False): + return + except Exception: + return + + mirrored_modules = ("torch._dynamo.config", "torch._inductor.config") + + # config.patch(...) and config.load_config(...) also assign via __setattr__, but + # their writes are thread-local by design; a per-thread depth counter marks them + # so they are not mirrored into the process-global default. + import threading + + _scoped_depth = threading.local() + + def _in_scoped_write(): + return getattr(_scoped_depth, "n", 0) > 0 + + def _bump(delta): + _scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta + + original_patch = ConfigModule.patch + if not getattr(original_patch, "__unsloth_patched__", False): + + @functools.wraps(original_patch) + def _patched_patch(self, *args, **kwargs): + ctx = original_patch(self, *args, **kwargs) + try: + cls = type(ctx) # patch() builds a fresh ConfigPatch class each call + if not getattr(cls, "__unsloth_patch_wrapped__", False): + _enter0, _exit0 = cls.__enter__, cls.__exit__ + + def _enter(s, _e = _enter0): + _bump(1) + try: + return _e(s) + finally: + _bump(-1) + + def _exit( + s, + *a, + _x = _exit0, + ): + _bump(1) + try: + return _x(s, *a) + finally: + _bump(-1) + + cls.__enter__, cls.__exit__ = _enter, _exit + cls.__unsloth_patch_wrapped__ = True + except Exception: + pass + return ctx + + _patched_patch.__unsloth_patched__ = True + ConfigModule.patch = _patched_patch + + # load_config restores a saved config by calling setattr per key (thread-local). + original_load_config = getattr(ConfigModule, "load_config", None) + if callable(original_load_config) and not getattr( + original_load_config, "__unsloth_patched__", False + ): + + @functools.wraps(original_load_config) + def _patched_load_config(self, *args, **kwargs): + _bump(1) + try: + return original_load_config(self, *args, **kwargs) + finally: + _bump(-1) + + _patched_load_config.__unsloth_patched__ = True + ConfigModule.load_config = _patched_load_config + + @functools.wraps(original_setattr) + def _patched_setattr(self, name, value): + original_setattr(self, name, value) + if _in_scoped_write(): + return # transient patch / load_config write: keep it thread-local + # Aliases (cache_size_limit -> recompile_limit) re-enter with the real name. + if self.__dict__.get("__name__", None) in mirrored_modules: + try: + entry = self.__dict__["_config"].get(name, None) + if entry is not None and entry.alias is None: + entry.default = value + except Exception: + pass + + _patched_setattr.__unsloth_patched__ = True + ConfigModule.__setattr__ = _patched_setattr + + # No replay of existing overrides: unsloth installs this before it sets any + # dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying + # would also bake a still-active config.patch override into the global default. + logger.info( + "Unsloth: Patched torch config modules so dynamo/inductor settings " + "(e.g. recompile_limit) apply across threads on torch >= 2.12." + ) + + def patch_trunc_normal_precision_issue(): """ Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32. @@ -1323,8 +1452,7 @@ def fix_vllm_pdl_blackwell(): if patched: logger.info( - f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - " - f"patched: {', '.join(patched)}" + f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}" ) else: # Just set the env var - vLLM might be an older version without supports_pdl From b509d47dd7427a1ba9ff1c80d1ca64fb9889bddf Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:26:36 -0700 Subject: [PATCH 044/402] Silence torch._check_is_size FutureWarning and shim it if torch removes it (#7023) * Silence torch._check_is_size FutureWarning and shim it if torch removes it bitsandbytes 4-bit dequant calls torch._check_is_size, which torch deprecated with a FutureWarning ("Use _check(i >= 0) instead") that prints on every bnb-4bit load. Silence that warning in suppress_cuda_printf, and add fix_torch_check_is_size so a future torch that removes _check_is_size gets it shimmed to _check(i >= 0) (honoring the max bound) and bitsandbytes keeps working. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten fix_torch_check_is_size docstring Lead with what the shim does and drop the redundant line; two lines instead of three, same intent. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- unsloth/_gpu_init.py | 3 +++ unsloth/import_fixes.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/unsloth/_gpu_init.py b/unsloth/_gpu_init.py index e6178e60f3..984057e9f7 100644 --- a/unsloth/_gpu_init.py +++ b/unsloth/_gpu_init.py @@ -26,6 +26,7 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules] # Fix some issues before importing other packages from .import_fixes import ( fix_message_factory_issue, + fix_torch_check_is_size, check_fbgemm_gpu_version, disable_broken_causal_conv1d, disable_broken_vllm, @@ -72,6 +73,7 @@ fix_bitsandbytes_rocm_arch_detection() disable_broken_causal_conv1d() disable_broken_vllm() fix_message_factory_issue() +fix_torch_check_is_size() check_fbgemm_gpu_version() torchvision_compatibility_check() fix_diffusers_warnings() @@ -81,6 +83,7 @@ del fix_bitsandbytes_rocm_arch_detection del disable_broken_causal_conv1d del disable_broken_vllm del fix_message_factory_issue +del fix_torch_check_is_size del check_fbgemm_gpu_version del torchvision_compatibility_check del fix_diffusers_warnings diff --git a/unsloth/import_fixes.py b/unsloth/import_fixes.py index c5300ed4d0..09de248c7b 100644 --- a/unsloth/import_fixes.py +++ b/unsloth/import_fixes.py @@ -172,6 +172,10 @@ if not UNSLOTH_ENABLE_LOGGING: # Deprecation warnings from torchao warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated") warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated") + # torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant) + warnings.filterwarnings( + "ignore", message = r"_check_is_size will be removed", category = FutureWarning + ) # TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752) warnings.filterwarnings( @@ -253,6 +257,30 @@ if not UNSLOTH_ENABLE_LOGGING: ) +def fix_torch_check_is_size(): + """Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit + dequant calls it). The FutureWarning is silenced in suppress_cuda_printf.""" + try: + import torch + + if hasattr(torch, "_check_is_size"): + return + + def _check_is_size( + i, + message = None, + *, + max = None, + ): + torch._check(i >= 0, message) + if max is not None: + torch._check(i <= max, message) + + torch._check_is_size = _check_is_size + except Exception: + return + + # Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype' # MUST do this at the start primarily due to tensorflow causing issues def fix_message_factory_issue(): From c1e06e9ddfb53a9d40e1fb182030aded94f4b4bb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 02:47:59 -0700 Subject: [PATCH 045/402] unsloth start: add --persist to keep and reopen agent sessions (#7014) * unsloth start: add --resume to persist and reopen agent sessions `unsloth start ` launches a coding agent whose home is a throwaway temp dir wiped on exit, so codex/openclaw/hermes/pi (which relocate their whole home there) cannot resume a conversation after you quit. opencode and claude keep their session data in a fixed user dir, so they already resume. Add an opt-in --resume/--no-resume flag: it routes the launch to the stable Unsloth agents dir (the same one --no-launch already uses) so the session survives the exit, never touching the user's own ~/.. A bare --resume also reopens the last conversation via the agent's native flag (codex `resume --last`, opencode/claude/pi `--continue`). The default is unchanged: a plain launch still uses a temp dir and persists nothing. Add a dispatch-only `resume` job to the Local Agent Guides CI that drives the real launch path and asserts the split: codex/pi are wiped without --resume and persist with it, while opencode/claude persist either way. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * unsloth start: rename --resume to --persist The session flag collided with agents' own resume flags. `unsloth start claude --resume ` used to forward `--resume ` straight to Claude (which keeps its history in ~/.claude regardless), so a boolean --resume on unsloth start would have swallowed the session id and turned it into a stray prompt. Name the persistence flag --persist instead, so every agent's native resume flag (claude --resume , codex resume, opencode --continue, ...) still passes through untouched. Behavior is otherwise identical: --persist keeps a launched agent's session under the Unsloth agents dir, and a bare --persist reopens the last conversation. Add a regression test that `--resume ` passes through verbatim, and in the CI resume experiment skip the redundant second pass for opencode/claude (they persist either way, and a second CPU turn only risks a timeout). * unsloth start: correct --persist help and drop the buggy auto-resume Reword the --persist help to be accurate: claude and opencode keep sessions in the user's own stores and resume regardless, so --persist only stabilizes the otherwise-ephemeral relocated home of codex/openclaw/hermes/pi. Drop the bare-launch auto-append of native resume tokens: it errored on a first launch with no prior session, and was inconsistent between launch and no-launch. --persist now only keeps the session dir; resume via the agent's own command (e.g. `unsloth start codex --persist resume`), which now finds it. In the CI resume experiment, fail the pass when the launched turn exits non-zero, so a write-then-error is not misread as PERSISTED. --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .github/scripts/agent-guides-drive.sh | 148 +++++++++++++++++ .github/workflows/local-agent-guides-ci.yml | 170 ++++++++++++++++++++ unsloth_cli/commands/start.py | 53 ++++-- unsloth_cli/tests/test_start.py | 142 ++++++++++++++++ 4 files changed, 502 insertions(+), 11 deletions(-) diff --git a/.github/scripts/agent-guides-drive.sh b/.github/scripts/agent-guides-drive.sh index defdb498c7..f4189a159e 100755 --- a/.github/scripts/agent-guides-drive.sh +++ b/.github/scripts/agent-guides-drive.sh @@ -527,6 +527,154 @@ case "$MODE" in echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)" ;; + # ── resume: does a launched agent's session survive exit and resume? ──── + # Unlike the other modes, this drives the real LAUNCH path (`unsloth start + # ...`, the interactive default), not the --no-launch recipe. That + # path relocates each agent's home to a throwaway temp dir wiped on exit, so + # a session cannot be resumed -- unless --persist routes it to the stable + # Unsloth agents dir instead. We run one headless turn per pass and check + # whether the turn left a session in a persistent store (deterministic, no + # reliance on the model recalling anything), for a baseline pass and a + # --persist pass, and assert the expected split for this agent. + resume) + CODEWORD="PLATYPUS7" + T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK." + T2="What codeword did I ask you to remember? Reply with just that word." + WORK="$WORKDIR_BASE/${AGENT}-resume" + + # STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to. + # Read it from a --no-launch probe (which also writes the agent's config + # there). codex/pi relocate their whole home/HOME here; opencode/claude keep + # their session data in a fixed user dir, so STABLE_HOME stays empty for them. + parse_connect + case "$AGENT" in + codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;; + pi) STABLE_HOME="$(raw_env HOME)" ;; + *) STABLE_HOME="" ;; + esac + + # The persistent stores a session would land in if it were NOT wiped. We + # count files here before/after each turn; a positive delta means the + # session persisted (is resumable), zero means it went to a wiped temp dir. + resume_tracked_dirs() { + case "$AGENT" in + codex) printf '%s\n' "$HOME/.codex" ;; + opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;; + claude) printf '%s\n' "$HOME/.claude" ;; + pi) printf '%s\n' "$HOME/.pi" ;; + *) : ;; + esac + [ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME" + } + count_session_files() { + local total=0 d n + while IFS= read -r d; do + [ -n "$d" ] && [ -d "$d" ] || continue + n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n)) + done < <(resume_tracked_dirs) + echo "$total" + } + + # The headless first-turn subcommand per agent (mirrors file-edit's map), + # forwarded verbatim through the launch path as passthrough args. + set_t1_cmd() { + case "$AGENT" in + claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;; + codex) T1_CMD=(exec "$T1") ;; + opencode) T1_CMD=(run "$T1") ;; + pi) T1_CMD=(-p "$T1") ;; + *) guide_fail "resume mode does not cover agent '$AGENT'" ;; + esac + } + + # Run one headless turn through the launch path. $1=outfile, $2="" or + # "--persist", rest = the agent subcommand. --yolo auto-approves so no tool + # prompt can hang; --api-key attaches to the already-served CI model. + launch_turn() { + local out="$1" rflag="$2"; shift 2 + local flag=(); [ -n "$rflag" ] && flag=("$rflag") + run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \ + --api-key "$UNSLOTH_API_KEY" "$@" + local rc=$? + redact "$out" + return "$rc" + } + + # One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED + # from the session-store delta. Runs in the main shell (not a command + # substitution) so a hang's guide_fail actually fails the job and the + # progress lines reach the CI log. $1 = "" (baseline) or "--persist". + RESULT="" + run_pass() { + local rflag="$1" label="baseline" + [ -n "$rflag" ] && label="resume" + rm -rf "$WORK"; mkdir -p "$WORK" + set_t1_cmd + local out="$LOGS_DIR/${AGENT}-resume-${label}.txt" + local before after rc + before="$(count_session_files)" + pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK" + launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$? + popd >/dev/null || true + after="$(count_session_files)" + echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})" + # The turn must succeed for the delta to mean anything: an agent that writes a + # session file then errors would otherwise be misread as PERSISTED. Mirror the + # file-edit mode and fail the pass on a non-zero launch (the flagship codex recall + # below stays WARN-only, driven by its own launch_turn calls). + [ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \ + guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; } + if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi + } + + run_pass ""; BASELINE="$RESULT" + # Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix. + # opencode/claude persist either way, so the baseline already proves it and a + # second full CPU turn only risks a timeout; skip it for them. + case "$AGENT" in + codex|pi) run_pass "--persist"; RESUME="$RESULT" ;; + *) RESUME="n/a (persists either way)" ;; + esac + + # Expected: codex/pi relocate their whole home to the temp dir, so a plain + # launch is WIPED and only --persist PERSISTS. opencode/claude keep their + # session data in a fixed user dir, so the baseline already PERSISTS. + case "$AGENT" in + codex|pi) EXPECT_BASELINE="WIPED" ;; + opencode|claude) EXPECT_BASELINE="PERSISTED" ;; + esac + + echo "──────────────────────────────────────────────" + echo "[$AGENT] RESUME EXPERIMENT" + echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})" + echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}" + echo "──────────────────────────────────────────────" + + [ "$BASELINE" = "$EXPECT_BASELINE" ] \ + || guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}" + case "$AGENT" in + codex|pi) + [ "$RESUME" = "PERSISTED" ] \ + || guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;; + esac + + # Flagship behavioral proof (codex only, WARN-only): after a --persist plant, + # resume the session and check the model actually recalls the codeword. A + # miss is not a failure (the CI model is small); the mechanism gate above is + # the real assertion. + if [ "$AGENT" = "codex" ]; then + rm -rf "$WORK"; mkdir -p "$WORK" + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true + ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true + if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then + echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}" + else + echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed" + fi + fi + echo "[$AGENT] resume OK" + ;; + *) echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2 exit 2 diff --git a/.github/workflows/local-agent-guides-ci.yml b/.github/workflows/local-agent-guides-ci.yml index 47f75dc1ba..25796bd5cf 100644 --- a/.github/workflows/local-agent-guides-ci.yml +++ b/.github/workflows/local-agent-guides-ci.yml @@ -471,6 +471,176 @@ jobs: redacted-configs/ retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ + # Job: resume + # Does a conversation started with `unsloth start ` survive exit + # and resume? This drives the REAL launch path (not the --no-launch + # recipe the other jobs use). A plain launch relocates the agent home to + # a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the + # session to the stable Unsloth agents dir so it persists. opencode/claude + # keep their session data in a fixed user dir, so they persist either way. + # Dispatch-only: it is an end-to-end experiment, not a PR gate. + # ═════════════════════════════════════════════════════════════════════ + resume: + name: resume (${{ matrix.agent }}) + if: github.event_name == 'workflow_dispatch' + runs-on: ubuntu-latest + timeout-minutes: 60 + strategy: + fail-fast: false + matrix: + # codex/pi relocate their whole home (resume broken without --persist); + # opencode/claude keep session data in a fixed dir (resume already works). + # One agent from each class proves the split end to end; openclaw/hermes + # share codex's relocation mechanism and are covered by the unit tests. + agent: [codex, opencode, claude, pi] + env: + GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF + GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf + STUDIO_PORT: '18904' + steps: + - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + with: + persist-credentials: false + + - name: Linux deps for llama.cpp prebuilt + run: | + sudo apt-get update + sudo apt-get install -y --no-install-recommends \ + libcurl4-openssl-dev libssl-dev jq + + - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0 + with: + node-version: '22' + + - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0 + with: + python-version: '3.12' + cache: 'pip' + + - name: Restore GGUF model file + id: cache-gguf + uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + continue-on-error: true + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Download GGUF if cache miss + id: download-gguf + if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success' + env: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + python -m pip install --upgrade huggingface_hub + mkdir -p gguf-cache + bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache + + - name: Save GGUF model file + if: always() && steps.download-gguf.outcome == 'success' + uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5 + with: + path: gguf-cache + key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1 + + - name: Install Studio (--local, --no-torch) + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: | + mkdir -p logs + set -o pipefail + bash install.sh --local --no-torch 2>&1 | tee logs/install.log + + - name: Serve unsloth run --disable-tools (gemma-4-E4B) + run: | + unsloth studio reset-password + bash .github/scripts/serve-unsloth-run.sh \ + --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \ + --port "$STUDIO_PORT" --log-dir logs \ + --extra "--seed $UNSLOTH_SEED --temp 0" \ + --health-timeout 900 + + - name: Preflight the agent's API dialect (class-a isolation) + env: + AGENT: ${{ matrix.agent }} + run: | + set -uo pipefail + B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY" + preflight_fail() { + echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**."; + exit 1 + } + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \ + -H "Authorization: Bearer $K") || true + [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code" + case "$AGENT" in + claude) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code" + ;; + codex) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true + [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code" + ;; + *) + code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \ + -H "Authorization: Bearer $K" -H 'content-type: application/json' \ + --max-time 120 \ + -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true + [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code" + ;; + esac + echo "preflight OK for $AGENT" + + - name: Install agent CLI (class-b isolation) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-install.sh "$AGENT" + + - name: Resume experiment (launch path) + env: + AGENT: ${{ matrix.agent }} + run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT" + + - name: Collect server logs (debug) + if: always() + run: | + mkdir -p logs/studio-logs + cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true + if [ -n "${UNSLOTH_API_KEY:-}" ]; then + grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do + sed -i "s#${UNSLOTH_API_KEY}##g" "$f" 2>/dev/null || true + done + fi + + - name: Stop Studio + if: always() + run: | + if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then + kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true + fi + sleep 2 + ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true + + - name: Upload logs + if: always() + continue-on-error: true + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 + with: + name: resume-${{ matrix.agent }}-log + path: | + logs/ + agent-workdir/ + redacted-configs/ + retention-days: 7 + # ═════════════════════════════════════════════════════════════════════ # Job 3: prompt-cache # (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0 diff --git a/unsloth_cli/commands/start.py b/unsloth_cli/commands/start.py index a8665b9be3..48c0aca34b 100644 --- a/unsloth_cli/commands/start.py +++ b/unsloth_cli/commands/start.py @@ -133,6 +133,21 @@ _YOLO_OPTION = typer.Option( "flag/config. Any of the three spellings works for any agent." ), ) +_PERSIST_OPTION = typer.Option( + False, + "--persist/--no-persist", + help = ( + "Keep this agent's Unsloth-managed session dir so you can resume it later. " + "codex/openclaw/hermes/pi have their whole home relocated into an Unsloth dir " + "that is a throwaway temp dir (wiped on exit) by default; with --persist it " + "lives under the Unsloth agents dir and survives, so their own resume can reopen " + "it. claude and opencode keep sessions in your own stores (~/.claude, " + "~/.local/share/opencode), so they already resume regardless. To reopen a " + "session, pass the agent's own resume command through, e.g. " + "`unsloth start codex --persist resume` or `claude --resume `; those flow to " + "the agent unchanged." + ), +) # Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no # such flag (config only) and are handled in their config writers, so they are absent. @@ -1133,15 +1148,20 @@ def _agents_config_root() -> Path: @contextlib.contextmanager -def _session_config(agent: str, launch: bool): +def _session_config( + agent: str, + launch: bool, + persist: bool = False, +): """Yield a private directory for an agent's session config (never the user's own). - launch: an ephemeral temp dir removed after the agent process exits, so nothing - persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later - on this machine), reused across runs. Either way the user's real ~/. - config is left untouched. + launch (default): an ephemeral temp dir removed after the agent process exits, so + nothing persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run + later on this machine), reused across runs. persist (from --persist): use that same + stable dir even for a launch, so the agent's session survives the exit and can be + resumed next time. Either way the user's real ~/. config is left untouched. """ - if launch: + if launch and not persist: path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-")) try: yield path @@ -1453,6 +1473,7 @@ def claude( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point Claude Code at the running Studio server and start it.""" base, key, entry = _connect( @@ -1497,6 +1518,9 @@ def claude( # --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions. # IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a # sandbox is detected, and we don't want to falsely claim one on the user's host. + # claude keeps its history in ~/.claude/projects, which --settings/env never + # relocate, so a session already survives exit; resume it with `claude --continue` + # or `--resume ` passed through. command = [ "claude", "--model", @@ -1533,6 +1557,7 @@ def codex( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point OpenAI Codex at the running Studio server and start it.""" base, key, entry = _connect( @@ -1558,7 +1583,7 @@ def codex( *_yolo_command_flags("codex", yolo), *ctx.args, ] - with _session_config("codex", launch) as home: + with _session_config("codex", launch, persist = persist) as home: write_codex_config(base, entry, home) env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)} _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex") @@ -1576,6 +1601,7 @@ def openclaw( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point OpenClaw at the running Studio server and start it.""" base, key, entry = _connect( @@ -1601,7 +1627,7 @@ def openclaw( if os.name == "nt" else "curl -fsSL https://openclaw.ai/install.sh | bash" ) - with _session_config("openclaw", launch) as cfg: + with _session_config("openclaw", launch, persist = persist) as cfg: config_path = cfg / "openclaw.json" # key lives in the config, not the env; --yolo writes the exec policy here too. write_openclaw_config(base, key, entry, config_path, yolo = yolo) @@ -1622,6 +1648,7 @@ def opencode( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point OpenCode at the running Studio server and start it.""" base, key, entry = _connect( @@ -1645,7 +1672,9 @@ def opencode( command = ["opencode", "--model", opencode_model] else: command = ["opencode"] - with _session_config("opencode", launch) as cfg: + # opencode keeps sessions in ~/.local/share/opencode (never relocated), so resume + # already survives exit; reopen the last one by passing `opencode --continue` through. + with _session_config("opencode", launch, persist = persist) as cfg: config_path = cfg / "opencode.json" # OPENCODE_CONFIG is an overlay (loaded between the user's global and project # configs), so this adds the Unsloth provider/model for the session without @@ -1697,6 +1726,7 @@ def hermes( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point Hermes (Nous Research) at the running Studio server and start it.""" base, key, entry = _connect( @@ -1708,7 +1738,7 @@ def hermes( ) command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args] install_hint = _hermes_install_hint() - with _session_config("hermes", launch) as home: + with _session_config("hermes", launch, persist = persist) as home: # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state) # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session. write_hermes_config(base, entry, home / "config.yaml") @@ -1728,6 +1758,7 @@ def pi( tensor_parallel: bool = _TENSOR_PARALLEL_OPTION, serve: bool = _SERVE_OPTION, yolo: bool = _YOLO_OPTION, + persist: bool = _PERSIST_OPTION, ): """Point Pi (coding agent) at the running Studio server and start it.""" base, key, entry = _connect( @@ -1752,7 +1783,7 @@ def pi( # --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs # no install scripts), so accepting the prompt skips dependency lifecycle scripts. install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent" - with _session_config("pi", launch) as home: + with _session_config("pi", launch, persist = persist) as home: # Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers # it over $HOME/.pi/agent), so pin it at the session dir: an inherited # PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real diff --git a/unsloth_cli/tests/test_start.py b/unsloth_cli/tests/test_start.py index 87a295532e..065972b275 100644 --- a/unsloth_cli/tests/test_start.py +++ b/unsloth_cli/tests/test_start.py @@ -2548,3 +2548,145 @@ def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path with start._session_config("codex", launch = False) as home2: assert home2 == home assert (home2 / "sessions" / "live.sqlite").read_text() == "state" + + +# ── --persist: persist the agent session so it can be resumed ──────────────── +def test_session_config_persist_uses_stable_dir_and_survives(monkeypatch, tmp_path): + # --persist routes a launch to the stable Unsloth agents dir (the one --no-launch + # already uses) instead of a throwaway temp dir, and never wipes it on exit. + monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents") + with start._session_config("codex", launch = True, persist = True) as home: + assert home == tmp_path / "agents" / "codex" + (home / "marker").write_text("kept") + assert home.exists() + assert (home / "marker").read_text() == "kept" + + +def test_session_config_default_launch_is_ephemeral(): + # Default launch (no --persist) still uses a throwaway temp dir wiped on exit. + with start._session_config("codex", launch = True) as home: + assert home.exists() + assert "unsloth-codex-" in home.name + assert not home.exists() + + +# The temp-dir agents: --persist points each one's home/state env at the stable dir; +# without it, at an ephemeral temp path. opencode is handled separately (only its +# config overlay is relocated; its session data was never in the temp dir). +_RESUME_ENV_VAR = { + "codex": "CODEX_HOME", + "openclaw": "OPENCLAW_STATE_DIR", + "hermes": "HERMES_HOME", + "pi": "HOME", +} + + +def _capture_launch(monkeypatch, argv): + captured = {} + + def run( + command, + env = None, + **kwargs, + ): + captured["command"] = command + captured["env"] = env + return SimpleNamespace(returncode = 0) + + monkeypatch.setattr(start.subprocess, "run", run) + result = CliRunner().invoke(start.start_app, argv) + assert result.exit_code == 0, result.output + return captured + + +@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR)) +def test_resume_persists_agent_home_to_stable_dir(agent, fake_studio, tmp_path, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}") + captured = _capture_launch(monkeypatch, [agent, "--persist"]) + stable = tmp_path / "agents" / agent + assert captured["env"][_RESUME_ENV_VAR[agent]] == str(stable) + # The stable dir survives the agent exit, so the session can be resumed. + assert stable.exists() + + +@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR)) +def test_default_launch_home_is_ephemeral(agent, fake_studio, tmp_path, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}") + captured = _capture_launch(monkeypatch, [agent]) + home = captured["env"][_RESUME_ENV_VAR[agent]] + assert f"unsloth-{agent}-" in home + assert str(tmp_path / "agents") not in home + + +def test_resume_opencode_config_in_stable_dir(fake_studio, tmp_path, monkeypatch): + # opencode's session data lives in ~/.local/share/opencode (never relocated), so + # resume already survives exit; --persist also stabilizes its config overlay dir. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode") + captured = _capture_launch(monkeypatch, ["opencode", "--persist"]) + stable = tmp_path / "agents" / "opencode" + assert captured["env"]["OPENCODE_CONFIG"] == str(stable / "opencode.json") + assert stable.exists() + + +def test_persist_bare_codex_launch_has_no_resume_token(fake_studio, monkeypatch): + # A bare `--persist` only persists the session dir; it must NOT auto-append a native + # resume token, or the very first launch (no session yet) would send codex down its + # no-session error path. The user resumes explicitly: `unsloth start codex --persist resume`. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = _capture_launch(monkeypatch, ["codex", "--persist"]) + assert "resume" not in captured["command"] + # command[0] is the resolved executable path; assert the argv after it. + assert captured["command"][1:] == ["--oss", "--profile", start._CODEX_PROFILE] + + +def test_persist_bare_opencode_launch_has_no_resume_token(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode") + captured = _capture_launch(monkeypatch, ["opencode", "--persist"]) + assert "--continue" not in captured["command"] + assert captured["command"][1:] == ["--model", f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"] + + +def test_persist_bare_claude_launch_has_no_resume_token(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + captured = _capture_launch(monkeypatch, ["claude", "--persist"]) + assert "--continue" not in captured["command"] + assert captured["command"][1:] == ["--model", MODEL["id"]] + + +def test_resume_with_passthrough_does_not_auto_append(fake_studio, monkeypatch): + # When the caller drives their own subcommand, --persist only persists the dir; it + # must not inject a resume token that would collide with the user's command. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = _capture_launch(monkeypatch, ["codex", "--persist", "exec", "hello"]) + assert "resume" not in captured["command"] + assert captured["command"][-2:] == ["exec", "hello"] + + +def test_default_launch_has_no_resume_token(fake_studio, monkeypatch): + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex") + captured = _capture_launch(monkeypatch, ["codex"]) + assert "resume" not in captured["command"] + + +def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch): + # openclaw/hermes persist their session dir but have no non-interactive resume + # selector, so --persist must not append a token; their own picker resumes. + for agent in ("openclaw", "hermes"): + monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}") + captured = _capture_launch(monkeypatch, [agent, "--persist"]) + assert "resume" not in captured["command"] + assert "--continue" not in captured["command"] + + +def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch): + # The persistence flag is --persist, NOT --resume, so an agent's own + # `--resume ` (e.g. `unsloth start claude --resume `) still flows + # through to the agent verbatim and is not swallowed as a Studio option. + monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude") + monkeypatch.setattr(start, "_claude_flags", lambda: []) + captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"]) + assert captured["command"][-2:] == ["--resume", "some-session-guid"] + # Studio never auto-appends its own resume token when the user drives resume. + assert captured["command"].count("--resume") == 1 + assert "--continue" not in captured["command"] From eb775d320778bb496378344c061bd538e4d39ad9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 03:20:02 -0700 Subject: [PATCH 046/402] Studio /v1/messages: accept thinking and unknown content blocks (#7017) * Studio /v1/messages: accept thinking and unknown content blocks The Anthropic-compatible /v1/messages endpoint modeled a message's content as Union[str, list[{text|image|tool_use|tool_result}]], so any other block type made Pydantic reject the whole request with `messages.N.content.str: Input should be a valid string`. Resuming a Claude session commonly replays assistant turns that carry `thinking` (extended thinking) blocks, and sometimes a null content for a tool-only turn, both of which tripped this and returned a 400. Accept them: - Add a permissive AnthropicUnknownBlock fallback (any block whose type is not one of the four known ones), so thinking/redacted_thinking/provider-specific/ future blocks validate. A validator keeps known types on their typed models, so a malformed known block (e.g. a tool_use without id) still fails cleanly. - Coerce a null message (and tool_result) content to "" so the converter's `for block in content` stays safe. The converter already drops block types it does not translate, so a thinking block is not forwarded to the model. * Studio /v1/messages: keep user content validation strict Make the thinking/null leniency role-aware so it never silently drops real user input. Assistant turns (replayed history) still accept unknown/thinking blocks and coerce a null tool-only turn to empty. User turns keep the strict boundary: a null user content is rejected, and a content block the converter cannot translate is rejected instead of being dropped into an empty prompt. Also remove an empty file committed by accident. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio /v1/messages: coalesce resumed user turns and tighten content checks - The /v1/messages count and generation paths now coalesce the adjacent user turns that dropping an empty or null assistant turn can leave behind, so a strict GGUF chat template no longer 400s on non-alternating roles. - A user content block with a non-string type (list / dict) is rejected as a clean 400 instead of raising TypeError and escaping as a 500. - The assistant null-to-empty coercion only applies to an explicit null; an assistant turn that omits content entirely still fails required-field validation instead of being silently coerced to an empty string. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio /v1/messages: tighten comments --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/models/inference.py | 63 ++++++ studio/backend/routes/inference.py | 14 +- .../backend/tests/test_anthropic_messages.py | 184 ++++++++++++++++++ 3 files changed, 257 insertions(+), 4 deletions(-) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 0f27b695fe..53b0f14b09 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -1533,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel): tool_use_id: str content: Union[str, list] = "" + @field_validator("content", mode = "before") + @classmethod + def _coerce_null_content(cls, v): + # Some clients send null content for an empty tool result; the str|list + # union would 400 on it, so treat null as "". + return "" if v is None else v + + +# Block types the converter translates explicitly. Anything else (thinking / +# redacted_thinking, a provider block a resumed session replays, or a future type) +# is accepted as an unknown block and dropped by the converter, rather than 400-ing +# the whole request on strict validation. +_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"}) + + +class AnthropicUnknownBlock(BaseModel): + type: str + model_config = {"extra": "allow"} + + @field_validator("type") + @classmethod + def _only_unknown_types(cls, v): + # Known types parse as their typed models above (so a malformed known block + # still fails cleanly); this fallback only catches the rest. + if v in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError("known block type handled by its typed model") + return v + AnthropicContentBlock = Union[ AnthropicTextBlock, AnthropicImageBlock, AnthropicToolUseBlock, AnthropicToolResultBlock, + AnthropicUnknownBlock, ] @@ -1583,6 +1612,40 @@ class AnthropicMessage(BaseModel): role: Literal["user", "assistant"] content: Union[str, list[AnthropicContentBlock]] + @model_validator(mode = "before") + @classmethod + def _normalize_content(cls, data): + # Role-aware leniency that never silently drops real user input: + # - assistant: a resumed tool-only turn's null content -> "" (str|list would + # 400 on null; "" keeps the converter's `for block in content` safe). + # Unknown blocks (thinking / future types) validate via + # AnthropicUnknownBlock and are dropped by the converter. + # - user: keep strict. Null user content stays None so str|list rejects it + # (400) rather than forwarding an empty prompt; and reject block types the + # converter cannot translate, since it silently skips unknown user blocks + # -- a user turn made only of them would validate yet send no content + # (silent data loss). + if not isinstance(data, dict): + return data + content = data.get("content") + if data.get("role") == "assistant": + # Coerce only an explicit null (resumed tool-only turn). A missing + # content key stays malformed so the required-field check still 400s. + if "content" in data and content is None: + return {**data, "content": ""} + return data + if isinstance(content, list): + for block in content: + btype = ( + block.get("type") if isinstance(block, dict) else getattr(block, "type", None) + ) + # Guard the value: a non-string type is unsupported too, and a + # membership test on an unhashable value would raise TypeError + # (escaping as a 500 instead of a clean 400). + if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES: + raise ValueError(f"unsupported content block type {btype!r} in a user message") + return data + class AnthropicTool(BaseModel): # Client tools have input_schema; server tools may only have type/name. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index d9901a5b2e..4bb9ce655e 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -10088,8 +10088,11 @@ async def anthropic_count_tokens( # Apply the same sanitization /messages does before generation, so the count # matches the prompt the real request would build (otherwise empty-assistant # sentinels / synthetic tool history inflate the count or hit the fallback). - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) openai_tools = anthropic_tools_to_openai(payload.tools or []) or None @@ -10217,8 +10220,11 @@ async def anthropic_messages( # builders apply the same strip; without it an Anthropic /v1/messages caller # replaying a prior provider-side tool_use forwards fake builtin tool # history to a backend with no matching function declarations. - openai_messages = _strip_provider_synthetic_tool_history( - _drop_empty_assistant_sentinels(openai_messages) + # Coalesce adjacent user turns left behind by dropping an empty / null assistant + # turn, so a strict GGUF chat template does not 400 on non-alternating roles + # (mirrors the GGUF chat path); a no-op for already-alternating histories. + openai_messages = _coalesce_consecutive_user_turns( + _strip_provider_synthetic_tool_history(_drop_empty_assistant_sentinels(openai_messages)) ) # Enforce vision guard + re-encode embedded images to PNG so the Anthropic diff --git a/studio/backend/tests/test_anthropic_messages.py b/studio/backend/tests/test_anthropic_messages.py index 0c6550a3bb..54454f6563 100644 --- a/studio/backend/tests/test_anthropic_messages.py +++ b/studio/backend/tests/test_anthropic_messages.py @@ -1770,3 +1770,187 @@ class TestAnthropicMessagesToolRouting: _drive(anthropic_messages(payload, request = None, current_subject = "t")) assert backend.calls[0][0] == "plain" + + +def test_resumed_session_thinking_and_null_content_do_not_400(): + # A resumed session replays assistant turns with `thinking` (and sometimes null) + # content. Those must be accepted (thinking dropped by the converter), not 400ed. + from pydantic import ValidationError + + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "secret reasoning", "signature": "s"}, + {"type": "text", "text": "the answer"}, + {"type": "tool_use", "id": "t1", "name": "f", "input": {}}, + ], + }, + {"role": "assistant", "content": None}, # tool-only turn serialized as null + ], + ) + # Known blocks still parse as their typed models; only the unknown one is loose. + assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock" + assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock" + assert req.messages[2].content == "" # null coerced + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content")) + assert assistant["content"] == "the answer" + assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded + + # A malformed KNOWN block still fails cleanly instead of being swallowed. + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}], + ) + + +def test_user_null_content_rejected(): + # The null->"" leniency is assistant-only; a null user content must be rejected + # at the boundary, not coerced into an empty prompt and forwarded to the model. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": None}], + ) + + +def test_user_unknown_block_rejected_not_silently_dropped(): + # The converter skips user blocks it cannot translate, so a user turn whose only + # block is unknown would validate yet forward no content. Reject at the boundary + # to avoid that silent data loss (the assistant fallback is unaffected). + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "document", "source": {}}]}, + ], + ) + + +def test_user_translatable_blocks_still_accepted(): + # text / image / tool_result are translatable, so a real user message built from + # them must still pass; the unknown-block guard only trips on other types. + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is this?"}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": "AA"}, + }, + {"type": "tool_result", "tool_use_id": "t1", "content": "ok"}, + ], + } + ], + ) + assert [type(b).__name__ for b in req.messages[0].content] == [ + "AnthropicTextBlock", + "AnthropicImageBlock", + "AnthropicToolResultBlock", + ] + + openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages]) + assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai) + + +def test_user_malformed_known_block_still_rejected(): + # The guard only allow-lists a user block's *type*; the union still validates its + # shape, so a known-but-malformed block (tool_result without tool_use_id) fails. + from pydantic import ValidationError + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": [{"type": "tool_result", "content": "x"}]}, + ], + ) + + +def test_user_content_block_non_string_type_rejected_cleanly(): + # A user block whose `type` is a non-string (unhashable list / dict, or a stray + # int) must fail as a clean validation error, not raise TypeError from the + # frozenset membership test and escape as a 500. + from pydantic import ValidationError + for bad_type in ([], {}, 5): + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "user", "content": [{"type": bad_type}]}], + ) + + +def test_assistant_missing_content_key_still_rejected(): + # The null -> "" leniency is only for an EXPLICIT null. An assistant message that + # omits content entirely stays malformed and must fail required-field validation. + from pydantic import ValidationError + + with pytest.raises(ValidationError): + AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [{"role": "assistant"}], + ) + # An explicit null is still accepted and coerced (regression guard). + req = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": None}, + ], + ) + assert req.messages[1].content == "" + + +def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch): + # user -> assistant(null) -> user is now accepted: the null assistant turn coerces + # to "" and is dropped. The route must then coalesce the two remaining user turns + # so a strict GGUF chat template does not 400 on non-alternating roles. + backend = _mock_backend(monkeypatch, context_length = 2048) + + class _Req: + state = SimpleNamespace() + url = SimpleNamespace(path = "/v1/messages") + method = "POST" + + async def is_disconnected(self): + return False + + payload = AnthropicMessagesRequest( + model = "x", + max_tokens = 16, + messages = [ + {"role": "user", "content": "first question"}, + {"role": "assistant", "content": None}, + {"role": "user", "content": "please continue"}, + ], + ) + + response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t")) + assert response.status_code == 200 + + [(_path, kwargs)] = backend.calls + user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"] + assert len(user_turns) == 1 # the two user turns were merged, not left adjacent + merged = user_turns[0]["content"] + if isinstance(merged, list): + merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict)) + assert "first question" in merged and "please continue" in merged From 350233512092fc6847b42050b7768f5f5e9a4578 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Thu, 9 Jul 2026 07:39:48 -0300 Subject: [PATCH 047/402] Studio: add Vulkan llama.cpp support (#5819) * Studio: add Vulkan llama.cpp support * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address gemini's feedback * Studio: move the Vulkan VRAM probe into a standalone script * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Improve Vulkan probe error reporting * Resolve llama-server symlink so Vulkan build is detected * Drop unreachable Vulkan fallback in GPU free-memory dispatcher * Skip the Intel GPU probe when NVIDIA or ROCm is present * Reserve host RAM headroom for Vulkan integrated GPUs * Add a `UNSLOTH_FORCE_VULKAN` environment variable * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Honor GGML_VK_VISIBLE_DEVICES, reserve discrete Vulkan VRAM headroom, and clear Intel GPU on --cpu-fallback * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Route Intel and forced-Vulkan hosts to the upstream Vulkan prebuilt, add arm64 Vulkan, keep Vulkan out of RAG auto-detect * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Clear the fork release pin when routing a Vulkan host to the upstream repo * Gate auto-Vulkan routing on no physical NVIDIA so hidden CUDA devices aren't used * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Pin Vulkan launches with --device Vulkan instead of the raw GGML_VK_VISIBLE_DEVICES index space * Let user --device override the Vulkan pin, and gate direct Vulkan asset picks on no physical NVIDIA * Update RAG auto-backend test mocks for the _resolve_auto binary and Vulkan probes * Keep the add_dll_directory handle alive through the Vulkan probe DLL loads * Revert RAG auto Vulkan guard, guard multi-backend Vulkan detection, and preserve forced Vulkan across updates * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use getattr for RTLD_GLOBAL in the Vulkan probe CDLL mode * Skip CUDA/ROCm APU and datacenter GPU tuning on Vulkan builds On a Vulkan llama.cpp build gpu_indices are ggml compact ordinals, not CUDA/ROCm physical ids, so _amd_apu_wants_unified_memory and _apply_datacenter_env were reading the wrong device. On a mixed AMD APU plus discrete GPU host that could raise a spurious system-RAM shortfall and block a valid discrete-GPU load. Gate all three call sites on not is_vulkan_backend; the Vulkan path already reserves iGPU host headroom and the backend ignores GGML_CUDA_* anyway. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tighten Vulkan-guard comment in load_model * Reduce comments in Vulkan support to be more succinct * Resolve shell-wrapper llama-server entrypoint to the real lib dir create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root when it cannot symlink into build/bin. _find_llama_server_binary returns that root entrypoint, but Path.resolve() does not follow a shell wrapper, so _llama_lib_dir returned the install root and _is_vulkan_backend missed libggml-vulkan.so -- silently skipping the Vulkan probe and --device pin on an otherwise valid Vulkan install. Follow the wrapper's exec target to build/bin. Regression test: test_shell_wrapper_entrypoint_resolves_to_real_lib_dir. * [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> Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com> Co-authored-by: danielhanchen --- .../backend/core/inference/_vulkan_probe.py | 110 ++++++++ studio/backend/core/inference/llama_cpp.py | 246 ++++++++++++++++-- .../tests/test_install_resolve_prebuilt.py | 170 ++++++++++++ studio/backend/tests/test_llama_cpp_update.py | 42 +++ .../tests/test_llama_cpp_vulkan_probe.py | 193 ++++++++++++++ studio/backend/utils/llama_cpp_update.py | 6 + studio/install_llama_prebuilt.py | 243 ++++++++++++++++- 7 files changed, 984 insertions(+), 26 deletions(-) create mode 100644 studio/backend/core/inference/_vulkan_probe.py create mode 100644 studio/backend/tests/test_llama_cpp_vulkan_probe.py diff --git a/studio/backend/core/inference/_vulkan_probe.py b/studio/backend/core/inference/_vulkan_probe.py new file mode 100644 index 0000000000..706346daad --- /dev/null +++ b/studio/backend/core/inference/_vulkan_probe.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Standalone free-VRAM probe for the bundled ggml Vulkan backend. + +Run in a short-lived subprocess (``python _vulkan_probe.py ``) so the +Vulkan instance never lives in the long-running backend process. Loads the +bundled ggml Vulkan backend from ```` and prints one +``\\t\\t\\t`` line per device to stdout. +Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi +order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU +sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses +it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm +fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM. + +Uses only the standard library so it stays runnable as a bare script. +""" + +import ctypes +import os +import sys + +# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ... +_GGML_BACKEND_DEVICE_TYPE_IGPU = 2 + + +def _igpu_flags(base, lib, count: int) -> list[bool]: + """Per-device integrated-GPU flags via ggml's backend registry. + + The Vulkan reg enumerates devices in the same order as + ``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device = + i``), so reg index == device ordinal. Returns all-False on any failure so + the reader never over-caps a discrete card. + """ + flags = [False] * count + try: + lib.ggml_backend_vk_reg.restype = ctypes.c_void_p + lib.ggml_backend_vk_reg.argtypes = [] + base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t + base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p] + base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p + base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t] + base.ggml_backend_dev_type.restype = ctypes.c_int + base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p] + + reg = lib.ggml_backend_vk_reg() + if not reg: + return flags + dev_count = base.ggml_backend_reg_dev_count(reg) + for i in range(min(count, dev_count)): + dev = base.ggml_backend_reg_dev_get(reg, i) + if dev: + flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU + except Exception: + # Best-effort: any failure degrades to "discrete" so the memory + # readings still get through instead of crashing the probe. + pass + return flags + + +def main() -> int: + if len(sys.argv) < 2: + return 0 + bindir = sys.argv[1] + + # Hold add_dll_directory's handle for the rest of main() (the documented + # idiom) so bindir stays on the search path while the sibling ggml DLLs + # resolve below. + _dll_dir = None + if sys.platform == "win32": + base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll" + try: + _dll_dir = os.add_dll_directory(bindir) + except Exception: + pass + else: + base_name, vk_name = "libggml-base.so", "libggml-vulkan.so" + + # RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr + # falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode). + _rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0) + try: + base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global) + lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global) + except OSError as e: + print(f"ggml-vulkan load failed: {e}", file = sys.stderr) + return 1 + + lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int + lib.ggml_backend_vk_get_device_count.argtypes = [] + lib.ggml_backend_vk_get_device_memory.restype = None + lib.ggml_backend_vk_get_device_memory.argtypes = [ + ctypes.c_int, + ctypes.POINTER(ctypes.c_size_t), + ctypes.POINTER(ctypes.c_size_t), + ] + + count = lib.ggml_backend_vk_get_device_count() + igpu = _igpu_flags(base, lib, count) + rows = [] + for i in range(count): + free, total = ctypes.c_size_t(0), ctypes.c_size_t(0) + lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total)) + rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value)) + sys.stdout.write("\n".join(rows)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index f61402aa5c..3ba9eff857 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -1436,6 +1436,50 @@ def _backfill_usage_from_timings(usage, timings): return out +def _vulkan_lib_filename() -> str: + return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so" + + +# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit +# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared +# system RAM, so hold back the same margin rather than inventing a larger one. +_IGPU_HOST_RESERVE_MIB = 1024 + + +def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int: + """Reserve host headroom on an integrated (shared-memory) Vulkan GPU. + + An iGPU's reported free "VRAM" is really free system RAM, so sizing + context/offload against all of it would push the host into swap or the OOM + killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from + ggml's device type, so a discrete card is never touched; only ever reduces. + """ + if not is_igpu: + return free_mib + return max(0, free_mib - _IGPU_HOST_RESERVE_MIB) + + +def _llama_lib_dir(binary: str) -> Path: + # The installer exposes llama-server as a top-level entrypoint into build/bin/, + # where the ggml backend libs live, so callers looking for sibling libs (Vulkan + # detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a + # symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to + # a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it + # cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's + # exec target too, so a wrapper-based install still finds build/bin. + resolved = Path(binary).resolve() + try: + with open(resolved, "rb") as _f: + _head = _f.read(256) + if _head.startswith(b"#!"): + _m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore")) + if _m: + return (resolved.parent / _m.group(1)).resolve().parent + except OSError: + pass + return resolved.parent + + def _is_external_link(path: Path) -> bool: """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink or a Windows directory junction / reparse point. Such a link resolves into @@ -2278,6 +2322,30 @@ class LlamaCppBackend: return total + @staticmethod + def _is_vulkan_backend(binary: Optional[str] = None) -> bool: + """True if the installed llama.cpp build is Vulkan-only. + + The official prebuilts are single-backend, so the Vulkan ggml lib next + to llama-server identifies a Vulkan build. Keeps the free-memory probe + and GPU pin in ggml's Vulkan device-index space. For a custom + multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer + to that backend (torch-usable, better-understood probe/pin). + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return False + lib_dir = _llama_lib_dir(binary) + if not (lib_dir / _vulkan_lib_filename()).is_file(): + return False + for _backend in ("cuda", "hip"): + sibling = ( + f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so" + ) + if (lib_dir / sibling).is_file(): + return False + return True + @staticmethod def _resolve_visible_physical_ids() -> Optional[list[int]]: """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on @@ -2440,11 +2508,42 @@ class LlamaCppBackend: return True @staticmethod - def _get_gpu_free_memory() -> list[tuple[int, int]]: + def _visible_devices_mask(env_name: str) -> Optional[set[int]]: + """Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset. + + ``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask + ("") yields an empty set (all devices hidden), distinct from an unset + var (None, no mask). Used by the nvidia-smi probe. + """ + raw = os.environ.get(env_name) + if raw is None: + return None + try: + return set(int(x.strip()) for x in raw.split(",") if x.strip()) + except ValueError: + return None + + @staticmethod + def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]: + """``--device Vulkan,...`` to pin a Vulkan launch to selected GPUs. + + The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory + reports and the registry names ``Vulkan``). Pin by that name, NOT via + GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw + vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices + and deduplicating ICDs), so a compact ordinal there could select a + different physical device or the CPU rasterizer. + """ + if not gpu_indices: + return [] + return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)] + + @staticmethod + def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]: """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by index; empty if no supported GPU is reachable. Thin wrapper over ``_get_gpu_memory`` for callers that only need free VRAM.""" - return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()] + return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)] @staticmethod def _apple_metal_memory_budget_bytes() -> int: @@ -2475,7 +2574,7 @@ class LlamaCppBackend: return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION) @staticmethod - def _get_gpu_memory() -> list[tuple[int, int, int]]: + def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]: """Query free AND total memory per GPU. Order: @@ -2487,9 +2586,18 @@ class LlamaCppBackend: probe returned [] on AMD) and NVIDIA hosts missing ``nvidia-smi`` from PATH. + On a Vulkan build the ggml Vulkan probe is authoritative, so the indices + are ggml's compact Vulkan ordinals (the space the pin selects via + ``--device Vulkan``). It reports ``total`` for discrete cards and 0 + for an iGPU (shared RAM) so the fit falls back to free*frac there. + Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm. + Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no - supported GPU is reachable. ``total`` lets the fit reserve absolute headroom. + supported GPU is reachable. """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if LlamaCppBackend._is_vulkan_backend(binary): + return LlamaCppBackend._get_gpu_free_memory_vulkan(binary) # ── NVIDIA via nvidia-smi ──────────────────────────────────── try: result = subprocess.run( @@ -2505,16 +2613,7 @@ class LlamaCppBackend: **_windows_hidden_subprocess_kwargs(), ) if result.returncode == 0: - allowed: Optional[set[int]] = None - cvd = os.environ.get("CUDA_VISIBLE_DEVICES") - if cvd is not None: - try: - # `if x.strip()` filters trailing-comma masks ("0,1,"). - # Empty mask (CVD="") yields an empty set -> all GPUs - # filtered out, per codebase convention. - allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip()) - except ValueError: - pass + allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES") gpus: list[tuple[int, int, int]] = [] for line in result.stdout.strip().splitlines(): parts = [p.strip() for p in line.split(",")] @@ -2579,6 +2678,91 @@ class LlamaCppBackend: logger.debug(f"torch GPU probe failed: {e}") return [] + @staticmethod + def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]: + """Query free (and total) VRAM per device via the bundled ggml Vulkan backend. + + Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance + in this process) and returns (device_index, free_mib, total_mib) sorted + by index. The index is ggml's compact Vulkan ordinal -- the one the + registry names ``Vulkan`` and load_model pins with ``--device``, + NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set + ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the + list already reflects it. iGPUs leave a host-RAM margin (see + ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass + their real total through. [] when no Vulkan build or device is reachable. + """ + binary = binary or LlamaCppBackend._find_llama_server_binary() + if not binary: + return [] + binary_dir = _llama_lib_dir(binary) + if not (binary_dir / _vulkan_lib_filename()).is_file(): + return [] + + env = child_env_without_native_path_secret() + # Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so + # the probe enumerates the same device list the launch will, named + # Vulkan0..N in the compact order reported here and pinned by that name + # via --device -- probe, mask, and pin stay in one index space. Do NOT + # filter the mask in Python: ggml parses the env var in raw + # vkEnumeratePhysicalDevices space while this probe reports the compact + # post-filter ordinal, so a Python filter would compare mismatched spaces. + if sys.platform != "win32": + # Let the loader resolve sibling ggml libs next to the binary. + existing_ld = env.get("LD_LIBRARY_PATH", "") + env["LD_LIBRARY_PATH"] = ( + f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir) + ) + probe_script = Path(__file__).with_name("_vulkan_probe.py") + try: + result = subprocess.run( + [sys.executable, str(probe_script), str(binary_dir)], + capture_output = True, + text = True, + timeout = 15, + env = env, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + logger.debug( + f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}" + ) + return [] + except Exception as e: + logger.debug(f"vulkan GPU probe failed: {e}") + return [] + + gpus: list[tuple[int, int, int]] = [] + for line in result.stdout.strip().splitlines(): + parts = line.split("\t") + if len(parts) != 4: + continue + try: + idx = int(parts[0]) + free_mib = int(parts[1]) // (1024 * 1024) + is_igpu = parts[2] == "1" + # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the + # fit stays on free*frac (the host reserve below is its + # headroom); a discrete card passes its real total through. + total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024) + except ValueError: + continue + capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu) + if capped < free_mib: + logger.info( + f"Vulkan device VK{idx} is an integrated GPU sharing system " + f"RAM; reserving {free_mib - capped}MiB host headroom " + f"({free_mib}->{capped}MiB usable)" + ) + gpus.append((idx, capped, total_mib)) + gpus.sort(key = lambda g: g[0]) + if gpus: + logger.info( + "Vulkan GPU memory detected: " + + ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus) + ) + return gpus + @staticmethod def _available_system_memory_mib() -> Optional[int]: """Available system RAM in MiB (psutil, then /proc/meminfo), or None if @@ -2807,7 +2991,8 @@ class LlamaCppBackend: def _llama_server_env_for_binary(binary: str) -> dict[str, str]: """Build a subprocess env that lets llama-server resolve native libs.""" env = child_env_without_native_path_secret() - binary_dir = str(Path(binary).parent) + # _llama_lib_dir resolves the llama-server symlink to the real build/bin. + binary_dir = str(_llama_lib_dir(binary)) if sys.platform == "win32": # Ordering: see _build_windows_path_dirs. #5106. @@ -5210,6 +5395,7 @@ class LlamaCppBackend: # Resolve llama-server now but defer a not-found error: a block-diffusion # GGUF uses the diffusion runner, and its arch is only known after the header. binary = self._find_llama_server_binary() + is_vulkan_backend = self._is_vulkan_backend(binary) # ── Phase 2: download (NO lock held, so cancel can proceed) ── # mtp_draft_path arrives set for local Gemma loads (detected @@ -5449,7 +5635,8 @@ class LlamaCppBackend: model_size = gguf_size + mmproj_size # 2-tuple gpus for existing logic + a total map for the absolute # per-GPU headroom (correct when the GPU is already partly used). - _gpu_mem = self._get_gpu_memory() + # Pass binary so a Vulkan build probes ggml's Vulkan ordinals. + _gpu_mem = self._get_gpu_memory(binary) gpus = [(idx, free) for idx, free, _t in _gpu_mem] total_by_idx = {idx: total for idx, _f, total in _gpu_mem} @@ -6222,7 +6409,12 @@ class LlamaCppBackend: # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an # oversize load the OS would otherwise kill mid-flight. Base model # only: an optional MTP drafter is dropped by the MTP-drop fallback. - if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices): + # CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals. + if ( + model_size is not None + and not is_vulkan_backend + and self._amd_apu_wants_unified_memory(gpu_indices) + ): _ram_msg = self._apu_ram_shortfall_message( model_size, self._available_system_memory_mib() ) @@ -6485,6 +6677,12 @@ class LlamaCppBackend: ", ".join(unsupported_cache_flags), ) + # Vulkan pins via --device (a cmd arg, unlike the env-based + # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's + # last-wins parsing lets a user --device override Studio's pick. + if is_vulkan_backend and gpu_indices is not None: + cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices) + # User pass-through args go last so llama.cpp's last-wins parsing # lets the user override Studio's auto-set flags. Already # validated by the route via validate_extra_args(). @@ -6536,23 +6734,25 @@ class LlamaCppBackend: env.setdefault("OMP_NUM_THREADS", "2") # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use - # shared system RAM. setdefault so a user value wins. - if self._amd_apu_wants_unified_memory(gpu_indices): + # shared system RAM. setdefault so a user value wins. Not on Vulkan + # (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids. + if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices): env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1") logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1") # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU). # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1. - if self._apply_datacenter_env(env, gpu_indices): + if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices): multi_gpu = self._effective_gpu_count(gpu_indices) > 1 logger.info( f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})" ) # Pin to selected GPU(s). On ROCm, narrowing only - # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full - # set, so set HIP_VISIBLE_DEVICES too. - if gpu_indices is not None: + # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so + # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device + # (above), not here. + if gpu_indices is not None and not is_vulkan_backend: pinned = ",".join(str(i) for i in gpu_indices) env["CUDA_VISIBLE_DEVICES"] = pinned try: diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index b825172a63..090d2932ea 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -55,6 +55,27 @@ def _host(**kw): return ilp.HostInfo(**base) +def test_force_cpu_clears_all_gpu_attributes_including_intel(): + # --cpu-fallback is the "select the CPU prebuilt even when a GPU is present" + # escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or + # the planner still prepends the Vulkan asset on an Intel-GPU host. + host = _host( + is_linux = True, + is_x86_64 = True, + has_usable_nvidia = True, + has_physical_nvidia = True, + has_rocm = True, + rocm_gfx_target = "gfx1100", + has_intel_gpu = True, + ) + forced = ilp._apply_host_overrides(host, force_cpu = True) + assert forced.has_usable_nvidia is False + assert forced.has_physical_nvidia is False + assert forced.has_rocm is False + assert forced.rocm_gfx_target is None + assert forced.has_intel_gpu is False + + def test_macos_upstream_pin_only_for_explicit_pre26_upstream(): pre26 = _host( system = "Darwin", @@ -313,3 +334,152 @@ def test_sm103_host_drops_cuda128_windows_build(): ) kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129]) assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name] + + +def _upstream_release(tag, asset_names): + return { + "tag_name": tag, + "assets": [ + {"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names + ], + } + + +def test_direct_upstream_arm64_intel_prefers_vulkan(): + # Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU + # second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset). + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-arm64" in kinds + assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz" + + +def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only(): + # A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical + # True, usable False) + an Intel iGPU must NOT get the Vulkan archive even + # when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES + # and could grab the reserved card. It falls through to the CPU asset. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-cpu"] + + +def test_direct_upstream_arm64_without_intel_is_cpu_only(): + host = _host(is_linux = True, is_arm64 = True, machine = "aarch64") + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + assert [a.install_kind for a in plan.attempts] == ["linux-arm64"] + + +def test_direct_upstream_x86_intel_prefers_vulkan(): + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + rel = _upstream_release( + "b9925", + ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"], + ) + plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest") + kinds = [a.install_kind for a in plan.attempts] + assert kinds[0] == "linux-vulkan", kinds + assert "linux-cpu" in kinds + + +def test_linux_vulkan_health_glob_matches_bare_cpu_lib(): + # The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU + # libs so a valid Vulkan install is not re-flagged unhealthy every check. + choice = ilp.AssetChoice( + repo = UPSTREAM, + tag = "b9925", + name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", + url = "https://example/x", + source_label = "upstream", + install_kind = "linux-vulkan", + ) + groups = ilp.runtime_payload_health_groups(choice) + assert ["libggml-cpu*.so*"] in groups + assert ["libggml-cpu-*.so*"] not in groups + + +def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin(): + # Routing fork -> upstream also drops the fork release pin, which is in a + # different tag namespace and would make the upstream resolver miss. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False) + assert repo == UPSTREAM + assert tag == "" + assert routed.has_intel_gpu is True + + +def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin(): + # A pin set WITH an explicit upstream repo is already on upstream -> kept. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + _routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False) + assert repo == UPSTREAM + assert tag == "b9596" + + +def test_route_to_vulkan_prebuilt_cpu_fallback_wins(): + # --cpu-fallback suppresses Vulkan routing even for an Intel host. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True) + assert repo == FORK + assert tag == "b9596-mix-abc" + assert routed is host + + +def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted(): + # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1): + # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or + # Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU. + host = _host( + is_linux = True, + is_x86_64 = True, + has_intel_gpu = True, + has_physical_nvidia = True, + has_usable_nvidia = False, + ) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted(): + # An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path. + host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True) + _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + + +def test_route_to_vulkan_prebuilt_non_intel_unchanged(): + host = _host(is_linux = True, is_x86_64 = True) + routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False) + assert repo == FORK + assert routed is host + + +def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): + # The --resolve-prebuilt probe must agree with the install path: an + # auto-detected Intel host resolves against upstream (Vulkan), not the fork. + monkeypatch.setattr( + ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True) + ) + seen, out = _run_resolve_capture_host(monkeypatch, capsys) + assert seen["repo"] == UPSTREAM + assert out["repo"] == UPSTREAM diff --git a/studio/backend/tests/test_llama_cpp_update.py b/studio/backend/tests/test_llama_cpp_update.py index 5138e90471..f405ebcbd1 100644 --- a/studio/backend/tests/test_llama_cpp_update.py +++ b/studio/backend/tests/test_llama_cpp_update.py @@ -448,6 +448,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path): assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5" +def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path): + # A Vulkan install (marker asset carries 'vulkan') must re-assert + # UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to + # CUDA/ROCm and silently replaces the Vulkan build. + install_dir = tmp_path / "llama.cpp" + binary = _write_install( + install_dir, + "b9493", + repo = "ggml-org/llama.cpp", + asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz", + ) + monkeypatch.setattr(upd, "_find_binary", lambda: binary) + monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py") + monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518") + + def _on_start(cmd): + _write_install( + install_dir, + "b9518", + repo = "ggml-org/llama.cpp", + asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz", + ) + + popen_kwargs: dict = {} + _patch_installer_popen( + monkeypatch, + lines = ["installed\n"], + on_start = _on_start, + captured_kwargs = popen_kwargs, + ) + + assert upd.start_update()["started"] is True + deadline = time.time() + 10 + while time.time() < deadline: + job = upd.get_update_status()["job"] + if job["state"] in ("success", "error"): + break + time.sleep(0.05) + assert job["state"] == "success", job + assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1" + + def test_start_update_reports_full_release_tag(monkeypatch, tmp_path): install_dir = tmp_path / "llama.cpp" binary = _write_install(install_dir, "b9595") diff --git a/studio/backend/tests/test_llama_cpp_vulkan_probe.py b/studio/backend/tests/test_llama_cpp_vulkan_probe.py new file mode 100644 index 0000000000..92aaab4873 --- /dev/null +++ b/studio/backend/tests/test_llama_cpp_vulkan_probe.py @@ -0,0 +1,193 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Vulkan free-VRAM reader regression tests on a synthetic probe output. + +Covers the post-probe handling in +``LlamaCppBackend._get_gpu_free_memory_vulkan``: + + * integrated GPUs (probe reports is_igpu=1) leave a flat per-device host + margin matching llama.cpp's --fit-target, so context auto-sizing can't + over-commit shared RAM, and report total 0 (shared RAM is not a budget), + * discrete GPUs (is_igpu=0) keep their free untouched and pass their real + total through so the fit can reserve absolute headroom, + * an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged + (ggml applies it), not stripped or filtered in Python -- the probe reports + ggml's compact ordinal, which load_model pins with ``--device Vulkan``. + +The ggml Vulkan library is never loaded: subprocess.run is mocked to emit +the tab-separated lines the real ``_vulkan_probe.py`` would print. +""" + +from __future__ import annotations + +import subprocess +import sys +import types as _types +from pathlib import Path +from unittest import mock + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import importlib as _importlib # noqa: E402 + + +def _maybe_stub(name: str, builder): + try: + _importlib.import_module(name) + except ImportError: + sys.modules[name] = builder() + + +def _build_loggers_stub(): + m = _types.ModuleType("loggers") + m.get_logger = lambda name: __import__("logging").getLogger(name) + return m + + +_maybe_stub("loggers", _build_loggers_stub) +_maybe_stub("structlog", lambda: _types.ModuleType("structlog")) + +from core.inference import llama_cpp as _llama_mod # noqa: E402 +from core.inference.llama_cpp import ( # noqa: E402 + LlamaCppBackend, + _llama_lib_dir, + _vulkan_lib_filename, +) + +MIB = 1024 * 1024 +GIB = 1024 * MIB + + +def _make_vulkan_install(tmp_path: Path) -> str: + """A binary whose sibling dir holds the Vulkan ggml lib, so the + reader's ``is_vulkan_backend`` sibling-file check passes.""" + bindir = tmp_path / "build" / "bin" + bindir.mkdir(parents = True) + binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server") + binary.write_bytes(b"stub") + (bindir / _vulkan_lib_filename()).write_bytes(b"stub") + return str(binary) + + +def _mock_probe(rows: list[str], captured_env: dict | None = None): + """Patch subprocess.run so the _vulkan_probe.py call returns ``rows`` + (already tab-formatted), recording the env it was launched with.""" + real_run = subprocess.run + + def fake_run(cmd, *args, **kwargs): + if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd): + if captured_env is not None: + captured_env.clear() + captured_env.update(kwargs.get("env") or {}) + return subprocess.CompletedProcess( + args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = "" + ) + return real_run(cmd, *args, **kwargs) + + return mock.patch("subprocess.run", side_effect = fake_run) + + +def _row( + idx: int, + free_bytes: int, + is_igpu: int, + total_bytes: int = 0, +) -> str: + return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}" + + +def test_integrated_gpu_leaves_host_margin(tmp_path): + binary = _make_vulkan_install(tmp_path) + # iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target). + # total stays 0: shared system RAM is not a VRAM budget for the fit. + rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus + + +def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path): + binary = _make_vulkan_install(tmp_path) + # 6 GiB free on a partially occupied 24 GiB card: free is untouched and the + # real total flows through so the fit reserves absolute headroom (CUDA/ROCm + # parity) instead of the looser free*frac budget. + rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus + + +def test_large_discrete_gpu_is_untouched(tmp_path): + binary = _make_vulkan_install(tmp_path) + # A 48 GiB discrete card stays untouched regardless of size; only the + # iGPU flag triggers the host margin, never a VRAM/RAM ratio. + rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)] + with _mock_probe(rows): + gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus + + +def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch): + # The mask is NOT stripped or filtered in Python: ggml parses it in raw + # physical-device space while this probe reports the compact post-filter + # ordinal, so mixing spaces would be wrong. It is passed through unchanged + # so ggml applies it to the same device list the launch will enumerate. + binary = _make_vulkan_install(tmp_path) + monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1") + captured: dict = {} + rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)] + with _mock_probe(rows, captured_env = captured): + LlamaCppBackend._get_gpu_free_memory_vulkan(binary) + assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured + + +def test_vulkan_pin_args_uses_device_names_not_env_mask(): + # Pin by compact device name via --device (the space the probe reports and + # the registry names), never by writing a compact ordinal into the raw + # GGML_VK_VISIBLE_DEVICES index space. + assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"] + assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"] + assert LlamaCppBackend._vulkan_pin_args(None) == [] + assert LlamaCppBackend._vulkan_pin_args([]) == [] + + +def test_vulkan_only_build_is_detected(tmp_path): + binary = _make_vulkan_install(tmp_path) + assert LlamaCppBackend._is_vulkan_backend(binary) is True + + +def test_multi_backend_build_is_not_vulkan_only(tmp_path): + # A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be + # treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan + # device; defer to the CUDA/HIP path instead. + binary = _make_vulkan_install(tmp_path) + cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so" + (_llama_lib_dir(binary) / cuda).write_bytes(b"stub") + assert LlamaCppBackend._is_vulkan_backend(binary) is False + + +@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX") +def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path): + # create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root + # when it cannot symlink; _find_llama_server_binary returns that root entrypoint, + # so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else + # _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently + # never engage on a valid Vulkan install. + import os + + binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib + bindir = Path(binary).parent + wrapper = tmp_path / "llama-server" + wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n') + os.chmod(wrapper, 0o755) + assert _llama_lib_dir(str(wrapper)) == bindir + assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/studio/backend/utils/llama_cpp_update.py b/studio/backend/utils/llama_cpp_update.py index c16ae91467..1bcbfbf95a 100644 --- a/studio/backend/utils/llama_cpp_update.py +++ b/studio/backend/utils/llama_cpp_update.py @@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path logger.info("llama update: installing", cmd = " ".join(cmd)) # Stream progress lines into job["progress"]. env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5") + # Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm + # box would otherwise re-route and silently replace the Vulkan build. + # Re-assert it via the same env flag setup uses (mirrors + # _rocm_install_args). + if asset and "vulkan" in asset.lower(): + env["UNSLOTH_FORCE_VULKAN"] = "1" proc = subprocess.Popen( cmd, stdout = subprocess.PIPE, diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 6c75e6c394..856ba71478 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -10,6 +10,7 @@ import argparse import atexit import errno import fnmatch +import glob import hashlib import json import os @@ -265,6 +266,7 @@ class HostInfo: has_physical_nvidia: bool has_usable_nvidia: bool has_rocm: bool = False + has_intel_gpu: bool = False rocm_gfx_target: str | None = None # (major, minor) from platform.mac_ver(); None off macOS or if unparseable. # Skips a macos prebuilt whose minimum-OS exceeds this host. @@ -1482,6 +1484,24 @@ def direct_upstream_release_plan( install_kind = "windows-hip", ) ) + # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. Gate + # on no PHYSICAL NVIDIA (not just no usable one): a host that hid NVIDIA + # via CUDA_VISIBLE_DEVICES must not reach Vulkan, which ignores that mask + # and could enumerate the reserved card. Falls through to CPU below. + elif host.has_intel_gpu and not host.has_physical_nvidia: + vulkan_asset = f"llama-{release_tag}-bin-win-vulkan-x64.zip" + vulkan_url = assets.get(vulkan_asset) + if vulkan_url: + attempts.append( + AssetChoice( + repo = repo, + tag = release_tag, + name = vulkan_asset, + url = vulkan_url, + source_label = "upstream", + install_kind = "windows-vulkan", + ) + ) cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip" cpu_url = assets.get(cpu_asset) if cpu_url: @@ -1545,6 +1565,23 @@ def direct_upstream_release_plan( # ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm # asset, so they fall through to the empty-attempts raise (HIP source # build) rather than silently getting a CPU binary on a GPU host. + # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. The + # elif already excludes usable NVIDIA and ROCm; also require no PHYSICAL + # NVIDIA so a CUDA-hidden card isn't reached through Vulkan (CPU below). + if host.has_intel_gpu and not host.has_physical_nvidia: + vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-x64.tar.gz" + vulkan_url = assets.get(vulkan_asset) + if vulkan_url: + attempts.append( + AssetChoice( + repo = repo, + tag = release_tag, + name = vulkan_asset, + url = vulkan_url, + source_label = "upstream", + install_kind = "linux-vulkan", + ) + ) asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz" asset_url = assets.get(asset_name) if asset_url: @@ -1564,6 +1601,23 @@ def direct_upstream_release_plan( # selector returned 0 attempts and the installer fell back to a # source build on every Linux ARM64 host (DGX Spark, Ampere # Altra, GitHub-hosted ubuntu-24.04-arm runners, etc.). + # Intel (or other non-NVIDIA/non-AMD) GPU: prefer the Vulkan prebuilt, + # mirroring the x86_64 branch. Upstream ships bin-ubuntu-vulkan-arm64. + # No physical NVIDIA: don't reach a CUDA-hidden card through Vulkan. + if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm: + vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-arm64.tar.gz" + vulkan_url = assets.get(vulkan_asset) + if vulkan_url: + attempts.append( + AssetChoice( + repo = repo, + tag = release_tag, + name = vulkan_asset, + url = vulkan_url, + source_label = "upstream", + install_kind = "linux-vulkan", + ) + ) asset_name = f"llama-{release_tag}-bin-ubuntu-arm64.tar.gz" asset_url = assets.get(asset_name) if asset_url: @@ -3075,6 +3129,40 @@ def detect_host() -> HostInfo: # Note: amdhip64.dll presence alone is NOT treated as GPU evidence # since the HIP SDK can be installed without an AMD GPU. + # Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs + # vendor id (0x8086); Windows queries the WMI video controller list. Only + # probed with no usable NVIDIA and no ROCm (matching the Vulkan branches), + # keeping the probe (notably the Windows powershell call) off that path. + has_intel_gpu = False + if not has_usable_nvidia and not has_rocm: + if is_linux: + for _vendor_file in glob.glob("/sys/class/drm/card*/device/vendor"): + try: + with open(_vendor_file) as _vf: + if _vf.read().strip().lower() == "0x8086": + has_intel_gpu = True + break + except OSError: + continue + elif is_windows: + _ps = shutil.which("powershell") or shutil.which("pwsh") + if _ps: + try: + _result = run_capture( + [ + _ps, + "-NoProfile", + "-Command", + "Get-CimInstance Win32_VideoController | " + "Select-Object -ExpandProperty Name", + ], + timeout = 15, + ) + if _result.returncode == 0 and "intel" in _result.stdout.lower(): + has_intel_gpu = True + except Exception: + pass + return HostInfo( system = system, machine = machine, @@ -3090,6 +3178,7 @@ def detect_host() -> HostInfo: has_physical_nvidia = has_physical_nvidia, has_usable_nvidia = has_usable_nvidia, has_rocm = has_rocm, + has_intel_gpu = has_intel_gpu, rocm_gfx_target = rocm_gfx_target, macos_version = macos_version, ) @@ -3126,6 +3215,7 @@ def _apply_host_overrides( has_physical_nvidia = False, has_rocm = False, rocm_gfx_target = None, + has_intel_gpu = False, ) gfx = _normalize_forwarded_gfx(override_rocm_gfx) if gfx: @@ -3866,6 +3956,23 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice "falling back to source build with HIP support" ) + # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. No + # physical NVIDIA (not just no usable one): a CUDA-hidden card must not + # be reached through Vulkan, which ignores CUDA_VISIBLE_DEVICES. + if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm: + vulkan_name = f"llama-{llama_tag}-bin-ubuntu-vulkan-x64.tar.gz" + if vulkan_name in upstream_assets: + log(f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}") + return AssetChoice( + repo = UPSTREAM_REPO, + tag = llama_tag, + name = vulkan_name, + url = upstream_assets[vulkan_name], + source_label = "upstream", + install_kind = "linux-vulkan", + ) + log("Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU") + upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz" if upstream_name not in upstream_assets: raise PrebuiltFallback("upstream Linux CPU asset was not found") @@ -3908,6 +4015,24 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice ) log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU") + # Intel (or other non-NVIDIA/non-AMD) GPU on Windows: use Vulkan. No + # physical NVIDIA so a CUDA-hidden card isn't reached through Vulkan. + if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm: + vulkan_name = f"llama-{llama_tag}-bin-win-vulkan-x64.zip" + if vulkan_name in upstream_assets: + log( + f"Intel GPU detected on Windows -- using upstream Vulkan prebuilt {vulkan_name}" + ) + return AssetChoice( + repo = UPSTREAM_REPO, + tag = llama_tag, + name = vulkan_name, + url = upstream_assets[vulkan_name], + source_label = "upstream", + install_kind = "windows-vulkan", + ) + log("Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU") + upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip" if upstream_name not in upstream_assets: raise PrebuiltFallback("upstream Windows CPU asset was not found") @@ -4503,6 +4628,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: "linux-arm64-cuda", "linux-rocm", "linux-arm64", + "linux-vulkan", }: return ["llama-server", "llama-quantize", "llama-diffusion-gemma-visual-server", "lib*.so*"] if choice.install_kind in {"macos-arm64", "macos-x64"}: @@ -4516,6 +4642,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]: "windows-cpu", "windows-cuda", "windows-hip", + "windows-vulkan", "windows-rocm", "windows-arm64", }: @@ -5731,8 +5858,10 @@ def validate_server( "linux-cuda", "linux-arm64-cuda", "linux-rocm", + "linux-vulkan", "windows-cuda", "windows-hip", + "windows-vulkan", "windows-rocm", "macos-arm64", } @@ -6354,6 +6483,20 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: ["libmtmd.so*"], ["libggml-hip.so*"], ] + if choice.install_kind == "linux-vulkan": + return [ + ["libllama-common.so*"], + ["libllama.so*"], + ["libggml.so*"], + ["libggml-base.so*"], + # Match the sibling globs (linux-cuda/-rocm): x64 bundles ship + # arch-suffixed libggml-cpu-.so, arm64 may ship a bare + # libggml-cpu.so; the '-' form missed the latter and re-flagged + # the install unhealthy on every check. + ["libggml-cpu*.so*"], + ["libmtmd.so*"], + ["libggml-vulkan.so*"], + ] if choice.install_kind in {"windows-cpu", "windows-arm64"}: return [["llama.dll"]] if choice.install_kind == "windows-cuda": @@ -6373,6 +6516,8 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]: return groups if choice.install_kind in {"windows-hip", "windows-rocm"}: return [["llama.dll"], ["*hip*.dll"]] + if choice.install_kind == "windows-vulkan": + return [["llama.dll"], ["ggml-vulkan.dll"]] return [] @@ -6654,6 +6799,89 @@ def validate_prebuilt_attempts( raise PrebuiltFallback("no prebuilt bundle passed validation") +def force_vulkan_requested() -> bool: + """Whether UNSLOTH_FORCE_VULKAN opts this host into the Vulkan llama.cpp + prebuilt instead of its detected CUDA/ROCm backend (e.g. so an AMD user can + run the Vulkan build for inference). Scoped to the llama.cpp backend; the + torch/training stack installs separately and still sees the real GPU. + """ + return os.environ.get("UNSLOTH_FORCE_VULKAN", "").strip().lower() in ( + "1", + "true", + "yes", + ) + + +def _vulkan_only_host(host: HostInfo) -> HostInfo: + """Rewrite ``host`` so the asset selectors take their Vulkan branch. + + That branch fires on ``has_intel_gpu and not nvidia and not rocm``, so clear + the CUDA/ROCm flags and raise the integrated-GPU flag. The synthetic flag + never leaves install planning -- it only routes the llama.cpp prebuilt + choice, not the torch/training stack. + """ + return dataclasses_replace( + host, + has_usable_nvidia = False, + has_physical_nvidia = False, + has_rocm = False, + has_intel_gpu = True, + ) + + +def _route_to_vulkan_prebuilt( + host: HostInfo, published_repo: str, published_release_tag: str, *, force_cpu: bool +) -> tuple[HostInfo, str, str]: + """Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt. + + The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes + from UPSTREAM_REPO. Two triggers route here, both suppressed under + --cpu-fallback (the explicit "give me CPU" last resort wins): + * UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend; + * an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose + of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset. + Applied by BOTH the install path and the --resolve-prebuilt probe so the + "is a prebuilt available" answer matches what actually gets installed. + + Returns the (possibly rewritten) host, repo, and release tag. + """ + forced = force_vulkan_requested() + # Gate auto-routing on no PHYSICAL NVIDIA, not merely no usable one: a mixed + # NVIDIA+Intel host that hides NVIDIA with CUDA_VISIBLE_DEVICES=""/-1 keeps + # has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores + # CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the + # reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides. + auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm + if force_cpu or not (forced or auto_intel): + return host, published_repo, published_release_tag + if host.is_macos: + if forced: + log( + "UNSLOTH_FORCE_VULKAN is set but ignored on macOS " + "(Metal is used; there is no Vulkan prebuilt)" + ) + return host, published_repo, published_release_tag + if forced: + log( + "UNSLOTH_FORCE_VULKAN is set; installing the upstream Vulkan " + "llama.cpp prebuilt instead of the detected GPU backend" + ) + # Forcing may override a detected NVIDIA/ROCm host, so normalize it to + # Vulkan-only; an auto-detected Intel host already is. + host = _vulkan_only_host(host) + else: + log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt") + # Swapping the fork for upstream invalidates a fork release pin: the two use + # different tag namespaces (fork b9596-mix- vs upstream b9596), so a + # pinned fork tag would make the upstream resolver query a nonexistent + # release and fall back to source. Drop it and let the upstream resolver + # pick by the requested llama tag. A pin already on an explicit upstream repo + # (repo unchanged here) is preserved. + if published_repo != UPSTREAM_REPO: + published_release_tag = "" + return host, UPSTREAM_REPO, published_release_tag + + def diffusion_visual_server_backfill_needed( install_dir: Path, host: HostInfo, choice: AssetChoice ) -> bool: @@ -6696,6 +6924,9 @@ def install_prebuilt( override_rocm_gfx = override_rocm_gfx, force_cpu = force_cpu, ) + host, published_repo, published_release_tag = _route_to_vulkan_prebuilt( + host, published_repo, published_release_tag, force_cpu = force_cpu + ) choice: AssetChoice | None = None try: with install_lock(install_lock_path(install_dir)): @@ -6708,7 +6939,9 @@ def install_prebuilt( f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install" ) # Single resolver: every fork host selects from the release manifest; - # an explicit ggml-org override selects by asset filename instead. + # an explicit ggml-org override selects by asset filename instead. A + # forced-Vulkan host already has published_repo pointed at + # UPSTREAM_REPO above, so the resolver takes the Vulkan asset branch. requested_tag, release_plans = resolve_simple_install_release_plans( llama_tag, host, @@ -6994,10 +7227,14 @@ def main() -> int: override_rocm_gfx = args.rocm_gfx, force_cpu = args.cpu_fallback, ) - repo = args.published_repo + # Same Vulkan routing the install path applies, so the probe's answer + # matches what would install (an Intel/forced-Vulkan host -> upstream). + host, repo, release_tag = _route_to_vulkan_prebuilt( + host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback + ) try: _requested, plans = resolve_simple_install_release_plans( - args.resolve_prebuilt, host, repo, args.published_release_tag or "" + args.resolve_prebuilt, host, repo, release_tag ) choice = plans[0].attempts[0] if plans and plans[0].attempts else None if choice is None: From 216a1fad33561ee4fcf24fd47811fbf721b46f29 Mon Sep 17 00:00:00 2001 From: alkinun Date: Thu, 9 Jul 2026 13:46:47 +0300 Subject: [PATCH 048/402] Fix Windows installer torch index override (#6972) * Fix Windows installer torch index override * Clear inherited uv index env vars for pinned installs in studio/setup.ps1 (#6898) * Harden setup.ps1 index-var clearing to truly remove vars (#6898) * Apply UV_DEFAULT_INDEX torch index fix to Linux/Mac install.sh (#6898) * Neutralize all uv index env vars for pinned torch installs (#6898) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- install.ps1 | 24 ++++++--- install.sh | 28 ++++++---- studio/setup.ps1 | 13 ++++- .../test_tokenizers_and_torch_constraint.py | 51 +++++++++++++++++++ tests/sh/test_mac_intel_compat.sh | 2 +- 5 files changed, 99 insertions(+), 19 deletions(-) diff --git a/install.ps1 b/install.ps1 index 696f4e613a..0797cd3868 100644 --- a/install.ps1 +++ b/install.ps1 @@ -469,6 +469,17 @@ function Install-UnslothStudio { param( [Parameter(Mandatory = $true)][ScriptBlock]$Command ) + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when the command pins an index, clear every uv index env var so + # it wins, then restore in finally. Other installs keep the user's mirror. + $savedUvIndex = $null + if ($Command.ToString() -match '--default-index') { + $savedUvIndex = @{} + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue + } + } $prevEap = $ErrorActionPreference $ErrorActionPreference = "Continue" try { @@ -488,6 +499,7 @@ function Install-UnslothStudio { return [int]$LASTEXITCODE } finally { $ErrorActionPreference = $prevEap + if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } } } } @@ -2200,7 +2212,7 @@ exit 0 # ABI-incompatible torchvision/torchaudio on AMD's per-arch index. $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec } if ($torchInstallExit -ne 0) { # Transient AMD-index failure: fall back to a CPU base so the install # still completes; Studio setup retries ROCm afterwards. @@ -2209,7 +2221,7 @@ exit 0 # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU # torch>= range, so without it uv would keep the ROCm build and only swap # the companions -- a mismatched venv the flavor-repair block won't fix. - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2223,7 +2235,7 @@ exit 0 } else { Write-TauriLog "STEP" "Installing PyTorch" substep "installing PyTorch ($TorchIndexUrl)..." - $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl } + $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl } if ($torchInstallExit -ne 0) { Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit) @@ -2306,7 +2318,7 @@ exit 0 # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx* - # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install + # is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops. if (-not $SkipTorch) { $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl @@ -2322,7 +2334,7 @@ exit 0 $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2331,7 +2343,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) diff --git a/install.sh b/install.sh index 0acc9ec0be..3f4ea92387 100755 --- a/install.sh +++ b/install.sh @@ -159,6 +159,12 @@ run_maybe_quiet() { run_install_cmd() { _label="$1" shift + # Installer-pinned index installs (torch) must beat an inherited uv mirror + # (#6898): when we pass --default-index, neutralize every uv index env var so + # the pinned index wins. Other installs keep the user's mirror. + case " $* " in + *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;; + esac if _is_verbose; then "$@" && return 0 _rc=$? @@ -2190,9 +2196,9 @@ _expected_torch_flavor_tag() { esac } -# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX / +# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX / # rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv -# resolves (torch + every transitive dep) via --index-url -- the same URLs the +# resolves (torch + every transitive dep) via --default-index -- the same URLs the # fresh-install paths above already use -- so a stale wheel is auto-repairable. # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall. _torch_index_repairable() { @@ -2744,7 +2750,7 @@ if [ "$_MIGRATED" = true ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2870,7 +2876,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" else substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..." # Pass explicit wheel URLs so the matched trio is @@ -2893,18 +2899,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN" run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm @@ -2964,7 +2970,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then substep "repairing ROCm torch (overwritten by dependency resolution)..." run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --force-reinstall fi ;; @@ -2999,14 +3005,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver") - # Repair when flavor is wrong AND the index is plain --index-url reinstallable + # Repair when flavor is wrong AND the index is plain --default-index reinstallable # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only. if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \ && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..." run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \ "$TORCH_CONSTRAINT" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" \ + --default-index "$TORCH_INDEX_URL" \ --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true) _installed_torch_tag="" @@ -3017,7 +3023,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN" substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN" substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN" - substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" + substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN" fi fi fi diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 07dcb17335..db01a1ecad 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2621,7 +2621,18 @@ function Fast-Install { param([Parameter(ValueFromRemainingArguments=$true)]$Args_) if ($UseUv) { $VenvPy = (Get-Command python).Source - $result = & uv pip install --python $VenvPy @Args_ 2>&1 + # An explicit --index-url must win. Inherited uv index env vars otherwise + # override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop + # them only for index-pinned installs; mirrors still apply elsewhere. + $saved = @{} + if (@($Args_) -contains '--index-url') { + foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') { + $saved[$n] = [Environment]::GetEnvironmentVariable($n) + Remove-Item "Env:$n" -ErrorAction SilentlyContinue + } + } + try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 } + finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } } if ($LASTEXITCODE -eq 0) { return } } & python -m pip install @Args_ 2>&1 diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index 7390d7be9b..4322f0c7d6 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -14,6 +14,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/ _REPO_ROOT = _TESTS_DIR.parent # unsloth/ _INSTALL_SH = _REPO_ROOT / "install.sh" _INSTALL_PS1 = _REPO_ROOT / "install.ps1" +_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1" _NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt" @@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged: assert '"torch>=2.4,<2.11.0"' in self._ps1 +class TestInstallPs1UvDefaultIndex: + """Installer-managed torch indexes must override inherited uv defaults.""" + + _ps1 = _read(_INSTALL_PS1) + + def test_torch_installs_use_default_index(self): + assert "--default-index $TorchIndexUrl" in self._ps1 + assert "--default-index $ROCmIndexUrl" in self._ps1 + + def test_torch_installs_do_not_use_deprecated_index_url(self): + assert "--index-url $TorchIndexUrl" not in self._ps1 + assert "--index-url $ROCmIndexUrl" not in self._ps1 + + def test_torch_installs_neutralize_all_uv_index_env_vars(self): + # Extra-index vars outrank --default-index, so pinned installs must clear them. + for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"): + assert var in self._ps1 + assert 'Remove-Item "Env:$n"' in self._ps1 + + +class TestSetupPs1FastInstallIndex: + """setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning.""" + + _ps1 = _read(_SETUP_PS1) + + def test_fast_install_clears_all_uv_index_env_vars(self): + for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"): + assert var in self._ps1 + # Must truly remove the vars (child sees no value), not set them empty. + assert 'Remove-Item "Env:$n"' in self._ps1 + + +class TestInstallShUvDefaultIndex: + """Linux/Mac installer torch indexes must override inherited uv defaults.""" + + _sh = _read(_INSTALL_SH) + + def test_torch_installs_use_default_index(self): + assert '--default-index "$TORCH_INDEX_URL"' in self._sh + + def test_torch_installs_do_not_use_deprecated_index_url(self): + assert '--index-url "$TORCH_INDEX_URL"' not in self._sh + + def test_torch_installs_neutralize_all_uv_index_env_vars(self): + # --default-index installs run with all uv index env vars unset via `env -u`. + assert ( + "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh + ) + + # Group 2 -- Shell snippet tests (bash subprocess, mocked python) class TestTorchConstraintShell: """Test the TORCH_CONSTRAINT block via bash with mocked python minor versions.""" diff --git a/tests/sh/test_mac_intel_compat.sh b/tests/sh/test_mac_intel_compat.sh index 3c3bbfaa5f..8a0ff4b641 100644 --- a/tests/sh/test_mac_intel_compat.sh +++ b/tests/sh/test_mac_intel_compat.sh @@ -312,7 +312,7 @@ if [ "$SKIP_TORCH" = true ]; then else echo "==> Installing PyTorch ($TORCH_INDEX_URL)..." uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \ - --index-url "$TORCH_INDEX_URL" + --default-index "$TORCH_INDEX_URL" fi TORCH_EOF From cd9d251f157bc8a014a68f4688b961344d5d02f6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 04:10:59 -0700 Subject: [PATCH 049/402] Fix fast inference crash on compressed-tensors FP8 models (#7025) * Fix fast_gemv crash on compressed-tensors FP8 models Loading a compressed-tensors FP8 checkpoint (for example unsloth/Llama-3.2-1B-Instruct-FP8-Block) with fast_inference=False and running a forward crashed with 'Parameter object has no attribute absmax' inside fast_gemv. A compressed-tensors CompressedLinear exposes an already dequantized bf16 weight at forward time while keeping a weight_scale Parameter. The quant state resolution in get_lora_parameters/get_lora_parameters_bias fell back to that weight_scale, so a bf16 weight was routed into the bitsandbytes fast_gemv/fast_dequantize path, which expects a bitsandbytes QuantState with an absmax attribute. Only fall back to weight_scale_inv/weight_scale when the weight is still fp8. A decompressed bf16 weight then resolves to no quant state and flows through the normal bf16 path, which already handles bias and the LoRA backward. Real fp8 and bitsandbytes 4bit weights are unchanged. * Skip the fast_gemv dispatch test before importing unsloth when bitsandbytes is absent --- tests/test_fast_gemv_dispatch.py | 63 ++++++++++++++++++++++++++++++++ unsloth/kernels/utils.py | 27 ++++++++++++-- 2 files changed, 86 insertions(+), 4 deletions(-) create mode 100644 tests/test_fast_gemv_dispatch.py diff --git a/tests/test_fast_gemv_dispatch.py b/tests/test_fast_gemv_dispatch.py new file mode 100644 index 0000000000..7758db2cd9 --- /dev/null +++ b/tests/test_fast_gemv_dispatch.py @@ -0,0 +1,63 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is +already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the +bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes. +""" + +from types import SimpleNamespace + +import pytest +import torch + +# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up +# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection +# errors instead of producing a skip. Any other import error still surfaces as a failure. +pytest.importorskip("bitsandbytes") + +import unsloth # noqa: F401 (sets UNSLOTH_IS_PRESENT before transformers) +from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES + +_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None + + +def _proj(weight, weight_scale = None): + proj = SimpleNamespace(weight = weight, bias = None, merged = False) + if weight_scale is not None: + proj.weight_scale = weight_scale + return proj + + +def test_bf16_weight_scale_not_used_as_quant_state(): + """A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None.""" + proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2)) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is None + + +def test_fp8_weight_keeps_scale(): + """An actual fp8 weight still resolves its weight_scale as the quant state.""" + if _FP8 is None: + pytest.skip("no float8 dtype in this torch build") + scale = torch.rand(2, 2) + proj = _proj(torch.randn(4, 4).to(_FP8), scale) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is scale + + +def test_plain_bf16_has_no_quant_state(): + proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16)) + W, W_quant = get_lora_parameters_bias(proj)[:2] + assert W_quant is None diff --git a/unsloth/kernels/utils.py b/unsloth/kernels/utils.py index 43ed198a4a..1b0b5ce12e 100644 --- a/unsloth/kernels/utils.py +++ b/unsloth/kernels/utils.py @@ -282,6 +282,21 @@ def QUANT_STATE(W): return getattr(W, "quant_state", None) +# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a +# quant state when the weight itself is still fp8. compressed-tensors layers expose an +# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around; +# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv / +# fast_dequantize path, which then reads a missing `absmax` and crashes. +_FP8_WEIGHT_DTYPES = tuple( + dtype + for dtype in ( + getattr(torch, "float8_e4m3fn", None), + getattr(torch, "float8_e5m2", None), + ) + if dtype is not None +) + + def get_lora_parameters(proj): """Return (weight, weight quant_state, lora A, lora B, lora scale). With QAT enabled, also fake-quantizes the base layer and lora weights. @@ -298,9 +313,11 @@ def get_lora_parameters(proj): if weight_fake_quantizer is not None: W = weight_fake_quantizer(W) - # Get quant state for 4bit or FP8 + # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the + # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer) + # must not carry a scale as its quant state or fast_gemv will crash on it. W_quant = getattr(W, "quant_state", None) - if W_quant is None: + if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES: W_quant = getattr(base_layer, "weight_scale_inv", None) if W_quant is None: W_quant = getattr(base_layer, "weight_scale", None) @@ -349,9 +366,11 @@ def get_lora_parameters_bias(proj): ) # (proj.base_layer if hasattr(proj, "base_layer") else proj) W = base_layer.weight - # Get quant state for 4bit or FP8 + # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the + # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer) + # must not carry a scale as its quant state or fast_gemv will crash on it. W_quant = getattr(W, "quant_state", None) - if W_quant is None: + if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES: W_quant = getattr(base_layer, "weight_scale_inv", None) if W_quant is None: W_quant = getattr(base_layer, "weight_scale", None) From 534c877d2136b47b0ceec25cc45900c7f7f15e6d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 04:20:41 -0700 Subject: [PATCH 050/402] Keep native RoPE scaling when extending context; carry rope_theta for linear (#7028) * Keep native RoPE scaling when extending context; carry rope_theta for linear When max_seq_length exceeds a model's native window, the loader overwrote the model's rope_scaling with linear scaling. For models that already ship a scaled RoPE (llama3/yarn/longrope) that is far worse for long context, and on transformers v5 the linear dict omitted rope_theta (v5 keeps it under rope_parameters), so the rotary base fell back to 10000 and broke past ~8K tokens. Keep the native scaling and just widen the window; only synthesize linear for plain-RoPE models, and carry rope_theta so v5 keeps the real base. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Only preserve native llama3 when extending context; keep linear fallback otherwise The patched attention constructor (patch_llama_rope_scaling) rebuilds only linear, llama3 and longrope and its longrope branch reads a top-level original_max_position_embeddings, so preserving yarn or a nested-only longrope config would raise during construction on transformers <= 4.47.1. Keep only llama3 native; yarn/longrope/other types fall back to the linear override, still carrying rope_theta. * Correct long-context extension comment to match llama3-only preservation --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/utils/test_rope_scaling_drift.py | 33 ++++++++++++++ unsloth/models/llama.py | 61 +++++++++++++++++--------- 2 files changed, 73 insertions(+), 21 deletions(-) diff --git a/tests/utils/test_rope_scaling_drift.py b/tests/utils/test_rope_scaling_drift.py index 7a738e236c..b2ec1e5a20 100644 --- a/tests/utils/test_rope_scaling_drift.py +++ b/tests/utils/test_rope_scaling_drift.py @@ -257,6 +257,39 @@ def test_recompute_helper_scales_on_cpu(): ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled." +def test_extended_rope_scaling_keeps_llama3_and_carries_theta(): + # Long-context extension keeps native llama3, but falls back to linear for every other + # type (the patched attention constructor only rebuilds linear/llama3/longrope), and the + # linear dict carries rope_theta so transformers v5 does not fall back to base 10000. + from types import SimpleNamespace + + from unsloth.models.llama import _extended_rope_scaling + + # llama3 model: keep native scaling, do not synthesize linear. + scaling, native = _extended_rope_scaling(_make_config(LLAMA3_ROPE_SCALING), 2.0) + assert ( + scaling is None and native == "llama3" + ), "must keep native llama3 scaling instead of overwriting it with linear." + + # yarn is not rebuildable by the patcher -> keep the safe linear fallback, not native. + yarn = SimpleNamespace(rope_scaling = {"rope_type": "yarn", "factor": 2.0}, rope_theta = 500000.0) + scaling, _ = _extended_rope_scaling(yarn, 2.0) + assert scaling == { + "type": "linear", + "factor": 2.0, + "rope_theta": 500000.0, + }, f"yarn must fall back to linear (patcher cannot rebuild it), got {scaling}." + + # plain RoPE with theta only under v5 rope_parameters: linear must carry rope_theta. + v5 = SimpleNamespace(rope_parameters = {"rope_type": "default", "rope_theta": 1000000.0}) + scaling, _ = _extended_rope_scaling(v5, 2.0) + assert scaling == { + "type": "linear", + "factor": 2.0, + "rope_theta": 1000000.0, + }, f"linear override dropped rope_theta on v5 (got {scaling}); base would fall back to 10000." + + def test_extended_rotary_reads_config_factor(): # LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8 # (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405). diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 1f43f61443..05523bc27b 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -1651,6 +1651,26 @@ def _rope_scaling_as_dict(rope_scaling): return {} +def _extended_rope_scaling(config, factor): + """RoPE scaling to extend a model past its native window. Keeps native llama3 as-is + (linear extension is far worse for long context); everything else gets linear. Returns + (scaling_or_None, type): None keeps llama3. The linear dict carries rope_theta so + transformers v5 (which stores it under rope_parameters) keeps the real base, not 10000. + Only llama3 is preserved because patch_llama_rope_scaling can only rebuild linear/llama3/ + longrope and its longrope branch needs a top-level original_max_position_embeddings.""" + existing = _rope_scaling_as_dict( + getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {} + ) + existing_type = existing.get("rope_type") or existing.get("type") + if existing_type == "llama3": + return None, existing_type + return { + "type": "linear", + "factor": factor, + "rope_theta": _get_rope_theta(config), + }, existing_type + + def _llama3_inv_freq_from_config( config, rope_scaling, @@ -2518,34 +2538,33 @@ class FastLlamaModel: max_seq_length = model_max_seq_length if (rope_scaling is None) and (max_seq_length > model_max_seq_length): - rope_scaling = max_seq_length / model_max_seq_length + factor = max_seq_length / model_max_seq_length if fast_inference: raise NotImplementedError( "Unsloth: Fast inference does not yet work with RoPE Scaling." ) - logger.warning_once( - f"Unsloth: {model_name} can only handle sequence lengths of at most " - f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of " - f"{round(rope_scaling, 3)}, it can be magically be extended to " - f"{max_seq_length}!" - ) - - # Warn RoPE scaling isn't allowed - if not has_rope_scaling: - raise RuntimeError( - f"However, {model_name} doesn't support RoPE Scaling!\n" - "Please file a feature request at https://github.com/unslothai/unsloth." + linear_scaling, native_type = _extended_rope_scaling(model_config, factor) + if linear_scaling is not None: + logger.warning_once( + f"Unsloth: {model_name} can only handle sequence lengths of at most " + f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of " + f"{round(factor, 3)}, it can be magically be extended to " + f"{max_seq_length}!" + ) + if not has_rope_scaling: + raise RuntimeError( + f"However, {model_name} doesn't support RoPE Scaling!\n" + "Please file a feature request at https://github.com/unslothai/unsloth." + ) + kwargs["rope_scaling"] = linear_scaling + else: + # Native llama3 scaling already handles long context; just widen the window. + logger.warning_once( + f"Unsloth: extending {model_name} to {max_seq_length} using its native " + f"{native_type} RoPE scaling." ) - - rope_scaling = { - "type": "linear", - "factor": rope_scaling, - } - - # Add to kwargs - kwargs["rope_scaling"] = rope_scaling from .loader_utils import ( check_and_disable_bitsandbytes_loading, From b5dca66cb1480b36ef738a4e62680e02cca2a65f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 04:52:30 -0700 Subject: [PATCH 051/402] scripts: refresh scan_packages allowlist baseline (#7032) * scripts: refresh scan_packages allowlist baseline Regenerate scripts/scan_packages_baseline.json against the current resolved dependency set so the blocking pip scan-packages gate matches what the scanner now finds. Refreshes evidence hashes for benign findings whose code shifted lines (unsloth-zoo mlx loader, gguf/mlx test /tmp fixtures) and adds two mainstream-library entries that were newly surfaced (torch inductor codecache base64+subprocess compile cache, torch testing common_utils socket import). Stale entries whose matching code changed and no longer triggers are dropped. All entries remain CRITICAL/HIGH findings manually judged benign; matched on (package, file, check, evidence_hash). * ci(security-audit): re-run scan when the allowlist baseline changes The security-audit pull_request trigger listed the scanners but not their allowlist baselines, so a baseline-only edit never re-ran the scan that consumes it. A refreshed baseline could therefore merge without CI confirming its evidence hashes match what the scanner finds. Add scan_packages_baseline.json and scan_npm_packages_baseline.json to the paths filter so baseline changes are validated on their own PR. --- .github/workflows/security-audit.yml | 6 +- scripts/scan_packages_baseline.json | 304 ++++++++++++--------------- 2 files changed, 140 insertions(+), 170 deletions(-) diff --git a/.github/workflows/security-audit.yml b/.github/workflows/security-audit.yml index 0ef2ad1e9d..1275d12216 100644 --- a/.github/workflows/security-audit.yml +++ b/.github/workflows/security-audit.yml @@ -2,8 +2,8 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. # Multi-language supply-chain audit. Triggers: -# - PRs touching any dependency manifest (Python / npm / Cargo) or -# this workflow file, +# - PRs touching any dependency manifest (Python / npm / Cargo), a +# scanner or its allowlist baseline, or this workflow file, # - push to main / pip, # - nightly @ 04:13 UTC so newly-published advisories surface even # when no PR opens, @@ -57,7 +57,9 @@ on: - 'studio/src-tauri/Cargo.lock' - 'pyproject.toml' - 'scripts/scan_packages.py' + - 'scripts/scan_packages_baseline.json' - 'scripts/scan_npm_packages.py' + - 'scripts/scan_npm_packages_baseline.json' - '.github/workflows/security-audit.yml' push: branches: [main, pip] diff --git a/scripts/scan_packages_baseline.json b/scripts/scan_packages_baseline.json index 1d34cfb66d..3582517d31 100644 --- a/scripts/scan_packages_baseline.json +++ b/scripts/scan_packages_baseline.json @@ -39,7 +39,7 @@ "file": "botocore/utils.py", "check": "Reads credential paths AND makes network calls", "severity": "CRITICAL", - "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3721: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", + "evidence": "Creds: L3551: CACHE_DIR = os.path.expanduser(os.path.join('~', '.aws', 'boto', 'cache')) | L3719: return os.path.expanduser(os.path.join('~', '.aws', 'login', 'cache'))\nNetwork: L32: from urllib.request import getproxies, proxy_bypass", "evidence_hash": "2d691bc373ab872aad23c744104596ba6d0d9f3b35aa101c7edbff4429b174c1" }, { @@ -55,23 +55,23 @@ "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", - "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" + "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", + "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" }, { "package": "datasets", "file": "datasets/utils/file_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L441: while True: sha256:ce92e38c17c524815e1f9055be77235028c1e68e41b45cbfe9c8f1b867a205da", - "evidence_hash": "cb36281d28a975d101121c0702ee05eeee470879520d39a8be552129333f514d" + "evidence": "L443: while True: sha256:feba37d77721aa658e1786d2e4b67de76fefe1ceeb3ce8529d361c5241778eea", + "evidence_hash": "2e458563dec752d0a9896c9685d368d9906867110db315ab751e3eb6ec63f51c" }, { "package": "diffusers", "file": "diffusers/utils/import_utils.py", "check": "Downloads and executes remote code", "severity": "CRITICAL", - "evidence": "L1015: return importlib.import_module(\".\" + module_name, self.__name__)", + "evidence": "L1052: return importlib.import_module(\".\" + module_name, self.__name__)", "evidence_hash": "e584ecfdb097d9482bb19cd3992813bc1a119cfd4c40af14748bafe22900d91e" }, { @@ -79,7 +79,7 @@ "file": "diffusers/utils/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L233: value = os.environ[key]\nNetwork: L688: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L709: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L728: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", + "evidence": "Env: L236: value = os.environ[key]\nNetwork: L691: response = requests.get(arry, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L712: response = requests.get(url, timeout=DIFFUSERS_REQUEST_TIMEOUT) | L731: image = PIL.Image.open(requests.get(image, stream=True, timeout=DIFFUSERS_REQUEST_TIMEOUT).raw)", "evidence_hash": "671190a6106c6ee9674e5e5942dc0940e1d2f8c78d5faf674413c2345b783fd9" }, { @@ -90,12 +90,20 @@ "evidence": "Archive: L317: a['TarFileType'] = tarfile.open(fileobj=_fileW,mode='w')\nNetwork: L330: x['SocketType'] = _socket = socket.socket()", "evidence_hash": "894862e547cf91b90cd6e4b495db3fb05b7490ef0d63de7e795a7e3d9447d850" }, + { + "package": "fastapi", + "file": "fastapi/routing.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", + "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" + }, { "package": "fastmcp-slim", "file": "fastmcp/cli/apps_dev.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L1340: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence": "Archive: L1353: with tarfile.open(fileobj=io.BytesIO(data), mode=\"r:gz\") as tar:\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", "evidence_hash": "73a7a72013e9f800627ea07e6dbc3beeb8c905a6a5480c8fd896f0063173d25c" }, { @@ -103,8 +111,8 @@ "file": "fastmcp/cli/apps_dev.py", "check": "Enumerates filesystem AND makes network calls", "severity": "CRITICAL", - "evidence": "FS: L624: history.replaceState(null, \"\", url); sha256:fd8dbfa8af4dea2ce43f4d441f3f81239de341b76a2eb0a33c446f6757ce5f43\nNetwork: L1291: with httpx.Client(timeout=30.0) as client: | L1305: with httpx.Client(timeout=30.0) as client: | L1335: with httpx.Client(timeout=30.0) as client: | L1537: client = httpx.AsyncClient(\nL1538: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1539: ) | L1701: async with httpx.AsyncClient(trust_env=False) as client: | L1769: with socket.socket(family, socket.SOCK_STREAM) as s:", - "evidence_hash": "6ada4a9111213bdee5ea24c70a72ec4acdc8ffe0de4a01fd9835bc261ccab8f8" + "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", + "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" }, { "package": "fonttools", @@ -132,19 +140,35 @@ }, { "package": "huggingface-hub", - "file": "huggingface_hub/hf_api.py", + "file": "huggingface_hub/_sandbox.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", - "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" + "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", + "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/_sandbox.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", + "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" }, { "package": "huggingface-hub", "file": "huggingface_hub/hf_api.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L4600: while True: sha256:f4a851312a1832efe1b435aa1275a82184e19cc3f47e2cd244373d56c11de272", - "evidence_hash": "dc8fcf44788e32f42d1cc2eb0e2deb55eb2dbf2c3a55909a7d503e450f45e602" + "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", + "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" + }, + { + "package": "huggingface-hub", + "file": "huggingface_hub/hf_api.py", + "check": "C2 polling/beaconing loop detected", + "severity": "CRITICAL", + "evidence": "L3746: while True: sha256:0c73ed1a7447120b112c063b14e720c6695bc11d00eb6b912cd0f10dc3e29b31", + "evidence_hash": "22f50b930e44146c5350bb99e6e6ebb09feea9bf1e899e407bedc4ffaf06721b" }, { "package": "huggingface-hub", @@ -159,8 +183,8 @@ "file": "huggingface_hub/utils/_http.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L443: while True: sha256:0ab4fed32d3af10f361963371f681923481377508a405b5d8770cef75f859168", - "evidence_hash": "1484f6b92f41c427ba8cbc7c4695a94975fea683dfa83aa510b4b0e982be4721" + "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", + "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" }, { "package": "huggingface-hub", @@ -218,6 +242,22 @@ "evidence": "L5: import socket sha256:915068303029fa5806199f256fb74504c65f253f9aee8ea23d8e384bb772b1c7", "evidence_hash": "30be130f165f418dfd37b144c5ae333de184b95f828ab8bd4010a67b84a5f814" }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", + "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" + }, + { + "package": "multiprocess", + "file": "multiprocess/tests/__init__.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", + "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" + }, { "package": "numba", "file": "numba/pycc/decorators.py", @@ -231,7 +271,7 @@ "file": "numba/tests/support.py", "check": "Reverse shell / bind shell pattern", "severity": "CRITICAL", - "evidence": "L1021: os.dup2(w, fd) | L1026: os.dup2(save, fd)", + "evidence": "L1016: os.dup2(w, fd) | L1021: os.dup2(save, fd)", "evidence_hash": "fea7aa03d48bf0f4386302fa444984c4f5dfc772cfec3f1df199fd33a52eec10" }, { @@ -495,16 +535,16 @@ "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", - "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" + "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", + "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" }, { "package": "scikit-learn", "file": "sklearn/datasets/_openml.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L100: while True: sha256:1f05a1b4fdd843b309634f583cb5e919866ef38ec5aa0b7d8a66ac8820655594", - "evidence_hash": "69597a64e5670a0f9a3c2aafc0bde4160f6170a9e2dc38f2c413cfa8d22ad193" + "evidence": "L100: while True: sha256:270363bb66980201e477f9b94886e4023f7a3d21b5ce026b7603a8c249a50c5b", + "evidence_hash": "53edbe07c312d459068d38e537b5114e65685ac3d4487b0423fa4542b5df20fe" }, { "package": "scikit-learn", @@ -642,6 +682,14 @@ "evidence": "Base64: L1211: content = base64.b64decode(data)\nSubprocess: L2692: subprocess.run(\nL2693: cmd.split(), capture_output=True, text=True, check=True\nL2694: ) | L2995: cmd_output = subprocess.run(\nL2996: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL2997: ) | L3707: out = subprocess.check_output(\nL3708: [\"ldd\", os.path.join(search, file)]\nL3709: ) | L3791: jobs.append(functools.partial(subprocess.check_call, cmd)) | L3876: subprocess.check_call(\nL3877: shlex.split(halide_cmd_gen.get_command_line())\nL3878: ) | L4336: subprocess.check_output(\nL4337: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4338: ) | L4591: output = subprocess.check_output(\nL4592: cmd_parts,\nL4593: stderr=subprocess.STDOUT,\nL4594: text=True,\nL4595: env=os.environ,\nL4596: )", "evidence_hash": "c09774087b702a6c5d6e2e85d9239c7c241ec938fbe9c0153e8f0b5c0710389b" }, + { + "package": "torch", + "file": "torch/_inductor/codecache.py", + "check": "base64 decode + subprocess execution (staged payload)", + "severity": "CRITICAL", + "evidence": "Base64: L1727: content = base64.b64decode(data)\nSubprocess: L3270: subprocess.run(\nL3271: cmd, capture_output=True, text=True, check=True\nL3272: ) | L3583: cmd_output = subprocess.run(\nL3584: (\"openssl\", \"sha512\", filename), capture_output=True, text=True\nL3585: ) | L4338: out = subprocess.check_output(\nL4339: [\"ldd\", os.path.join(search, file)]\nL4340: ) | L4422: jobs.append(functools.partial(subprocess.check_call, cmd)) | L4507: subprocess.check_call(\nL4508: shlex.split(halide_cmd_gen.get_command_line())\nL4509: ) | L4992: subprocess.check_output(\nL4993: cmd_parts, stderr=subprocess.STDOUT, env=os.environ\nL4994: ) | L5247: output = subprocess.check_output(\nL5248: cmd_parts,\nL5249: stderr=subprocess.STDOUT,\nL5250: text=True,\nL5251: env=os.environ,\nL5252: )", + "evidence_hash": "87f77b5f51cb84fe9950fdeeb90fe8710e1b863100e90b5e2cfb228a725bee06" + }, { "package": "torch", "file": "torch/ao/__init__.py", @@ -695,7 +743,7 @@ "file": "torch/testing/_internal/common_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L4770: env = os.environ.copy()\nNetwork: L4832: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4850: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", + "evidence": "Env: L4900: env = os.environ.copy()\nNetwork: L4962: with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: | L4980: with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:", "evidence_hash": "704a851b9d68c9b885b9e15538bd7e96f03875503b618fe6f126c4438edd7386" }, { @@ -706,6 +754,14 @@ "evidence": "L32: import socket sha256:89faaaa8bc908e02dad73fd59b2b481fa91189c84b39b556c2766e71d2783bf3", "evidence_hash": "3d23d77ace91812a07cb9508cf352185d154176e8e8c8b9b28fa92cdbcfe0d53" }, + { + "package": "torch", + "file": "torch/testing/_internal/common_utils.py", + "check": "Reverse shell / bind shell pattern", + "severity": "CRITICAL", + "evidence": "L32: import socket sha256:ba439cbf568b194872f1d974c02b0487e51f677b67e379400522d0992600bd2d", + "evidence_hash": "88e98b227573997f86eedea8e885a407b0dd549d46d4a3f0b840ec5aafe66865" + }, { "package": "torchvision", "file": "torchvision/datasets/utils.py", @@ -743,8 +799,8 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1663: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", - "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" + "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", + "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" }, { "package": "transformers", @@ -759,15 +815,15 @@ "file": "transformers/testing_utils.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L1577: while True: sha256:2c6152f9da685f728e58d39dfc1827bc794f52606f56983bf38b5c6d0857cd5b", - "evidence_hash": "cdada67f3327237f00838a6750a4908dfaf76b9ab30c1352495c340d4fbd15c9" + "evidence": "L1699: while True: sha256:969e911d30c37a279ad915fb8c3d2d0a3f5705a7eb82ae6e00687388b68bbe65", + "evidence_hash": "2aa8e94baa805d599720a16afee6f08976482e301333e619e6c343389498ad15" }, { "package": "transformers", "file": "transformers/testing_utils.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L284: value = os.environ[key] | L300: value = os.environ[key] | L2129: env = os.environ.copy() | L2251: for k in list(os.environ.keys()):\nNetwork: L2561: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", + "evidence": "Env: L288: value = os.environ[key] | L304: value = os.environ[key] | L2165: env = os.environ.copy() | L2287: for k in list(os.environ.keys()):\nNetwork: L2597: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:", "evidence_hash": "73ff16aee09cf163fb3a7a04dfa2cf610595bde2f19460a579397695f728e3f4" }, { @@ -799,16 +855,16 @@ "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", - "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" + "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", + "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" }, { "package": "trl", "file": "trl/extras/vllm_client.py", "check": "C2 polling/beaconing loop detected", "severity": "CRITICAL", - "evidence": "L152: while True: sha256:93e7d409e300af445376e6defbe2d0241aa19ecf63ed41b780fbb91c7d09856f", - "evidence_hash": "208838617172de61bca201d2a1bbeb5aa5aaa55feb1a1069cf39214673a7d6d1" + "evidence": "L146: while True: sha256:2beedc742e1f085eaa10fd3bc40be97d2331d21887ef1b9ccdfa2150a184edfe", + "evidence_hash": "1540dffaaa053780e953e04c11d9c6b9c74b91cb60f3e6d87451ba7fe7db46db" }, { "package": "trl", @@ -866,6 +922,14 @@ "evidence": "Crypto: L294: r\"|\\b(?:xprv|xpub|bc1|0x[a-fA-F0-9]{40})\\b\",\nNetwork: L53: import urllib.request | L1757: req = urllib.request.Request(url, headers = {\"Accept\": \"application/json\"}) | L1758: with urllib.request.urlopen(req, timeout = 30) as resp:", "evidence_hash": "278ff15b0b702d37d7f0b30a1e55a31bf2b11883685718a47478fbb5ce7f5212" }, + { + "package": "unsloth-zoo", + "file": "scripts/scan_packages.py", + "check": "Writes to /tmp and executes (staged dropper)", + "severity": "CRITICAL", + "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", + "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" + }, { "package": "unsloth-zoo", "file": "tests/security/fixtures/_build.py", @@ -919,16 +983,16 @@ "file": "tests/test_mlx_save_export_regressions.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L164: temporary_location=\"/tmp/ignored\", sha256:78837e80d48e872ef191aaacfe5e1c621a98a20df486a70a41d1a932d074a5b3", - "evidence_hash": "dd11376e664d0d7e7f4cc4baf57eacd4b7ae7b03222dce3912ce68b63dbfca1e" + "evidence": "L165: temporary_location=\"/tmp/ignored\", sha256:9f8502377b19666288b28399633dfc6740a64d0cb70ad1615e38b1269f94bf37", + "evidence_hash": "b7262d6e58f2ebad961dd3e64ca6c32bba356b5044d7a642d7dbd36a58cb6c81" }, { "package": "unsloth-zoo", "file": "tests/test_quantize_gguf_q2_k_l.py", "check": "Writes to /tmp and executes (staged dropper)", "severity": "CRITICAL", - "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:32532cadc357beee1009f4e86481bdbe60a0b7bf47f6bb022b05ec1b8e15aed0", - "evidence_hash": "49f5b67379de17178f21a9bc93b79d6b94a70ecbdd16de86574934aac30a071d" + "evidence": "L67: input_gguf=\"/tmp/in.gguf\", sha256:06789b55e8f31426c233f37ff7d3729cc9e1f61c0829abd2c00c39216c63c7ad", + "evidence_hash": "ad4913d9099eb9b70e09d6860b242eb5f48c67e46d9bf4ae35c1c38a267d753b" }, { "package": "unsloth-zoo", @@ -951,7 +1015,7 @@ "file": "unsloth_zoo/llama_cpp.py", "check": "Creates archive with sensitive data AND makes network calls", "severity": "CRITICAL", - "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence": "Archive: L938: with tarfile.open(archive_path, \"r:gz\") as archive:\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "b9f3b1652349fa8ef9ac2d1715978aca1e1632165851a00a2698dd47189e410c" }, { @@ -959,7 +1023,7 @@ "file": "unsloth_zoo/llama_cpp.py", "check": "Harvests environment variables/secrets AND makes network calls", "severity": "CRITICAL", - "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2862: check = requests.get(llama_cpp_chat_file, timeout = 5)", + "evidence": "Env: L125: keynames = \"\\n\" + \"\\n\".join(os.environ.keys()) | L683: token = os.environ.get(\"GH_TOKEN\") or os.environ.get(\"GITHUB_TOKEN\")\nNetwork: L691: response = requests.get(url, timeout = timeout, headers = headers, stream = stream) | L1699: response = requests.get(\nL1700: LLAMA_CPP_CONVERT_FILE, timeout = (10, 120)\nL1701: ) | L2873: check = requests.get(llama_cpp_chat_file, timeout = 5)", "evidence_hash": "9cd0b1bb59c7eb1d814d7636dfd167c34f265eb7c4521a9d88b2bdcfd535b926" }, { @@ -1002,6 +1066,14 @@ "evidence": "Obfusc: L87: __import__(name)\nExec: L735: exec(\"\"\"exec _code_ in _globs_, _locs_\"\"\")", "evidence_hash": "3cb7d8247dea7dd3d7b21ededc0181c58c50099aeb73c9138a286f3d1ad92d4f" }, + { + "package": "cffi", + "file": "cffi/_cffi_gen_src.py", + "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", + "severity": "HIGH", + "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", + "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" + }, { "package": "cffi", "file": "cffi/setuptools_ext.py", @@ -1127,7 +1199,7 @@ "file": "numba/tests/support.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L879: __import__(modname)\nExec: L813: eval(co, globs, ns)", + "evidence": "Obfusc: L874: __import__(modname)\nExec: L808: eval(co, globs, ns)", "evidence_hash": "649a7d750f903478243b0bcb9e8020521b505fc7fedc5b696ec01f4efc096109" }, { @@ -1159,7 +1231,7 @@ "file": "numba/tests/test_np_functions.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7106: exec(compile(funcstr, '', 'exec'), globals(), dct)", + "evidence": "Obfusc: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)\nExec: L7118: exec(compile(funcstr, '', 'exec'), globals(), dct)", "evidence_hash": "9e81164131d16056fb56ad3cd11b8d129d1ff4f5855031e8b501e0335d5c14ed" }, { @@ -1175,16 +1247,16 @@ "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", - "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" + "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", + "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" }, { "package": "numpy", "file": "numpy/testing/_private/utils.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L2788: original_trace = sys.gettrace() | L2790: sys.settrace(None) | L2793: sys.settrace(original_trace)\nSubprocess: L1486: output = subprocess.run(cmd, capture_output=True, text=True) | L2889: res = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True,\nL2890: errors=\"replace\", **kwargs)\nExec: L1352: exec(astr, dict) | L1640: exec(code, globs, locs)", - "evidence_hash": "9c6961817e5b1751e572dfe0858286703bb835870ecdfd6a7a9fdd8372a5dd2b" + "evidence": "Anti: L2777: original_trace = sys.gettrace() | L2779: sys.settrace(None) | L2782: sys.settrace(original_trace)\nSubprocess: L1478: output = subprocess.run(cmd, capture_output=True, text=True)\nExec: L1346: exec(astr, dict) | L1632: exec(code, globs, locs)", + "evidence_hash": "27468a6828101c6c026ae25aca8aa90ef485fd62b2c8f0967479edae9c965844" }, { "package": "numpy", @@ -1199,7 +1271,7 @@ "file": "PIL/Image.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3772: def eval(image: Image, *args: Callable[[int], float]) -> Image:", + "evidence": "Obfusc: L422: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), []) | L490: __import__(f\"{__spec__.parent}.{plugin}\", globals(), locals(), [])\nExec: L3776: def eval(image: Image, *args: Callable[[int], float]) -> Image:", "evidence_hash": "c2c1e7ae44e15862caf8de549d09db7b35e93282450f07ef61aaf5450a408c13" }, { @@ -1255,7 +1327,7 @@ "file": "setuptools/_distutils/compilers/C/base.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1286: __import__(module_name)\nExec: L1113: if lib_type not in eval(expected):", + "evidence": "Obfusc: L1287: __import__(module_name)\nExec: L1114: if lib_type not in eval(expected):", "evidence_hash": "368651e9818ed2d1bb009027d3bcfbf94ae30639c0882a6c2bddde97b8c4f1e5" }, { @@ -1271,7 +1343,7 @@ "file": "setuptools/tests/config/test_pyprojecttoml.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L364: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", + "evidence": "Obfusc: L387: \"setup.py\": \"__import__('setuptools').setup(include_package_data=False)\",\nExec: L98: \"__main__.py\": \"def exec(): print('hello')\",", "evidence_hash": "067d41014f72a61d8b4adf25f3659d1f66a0e909f732223f48837aa7684df4e6" }, { @@ -1279,7 +1351,7 @@ "file": "setuptools/tests/test_editable_install.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L449: exec(finder, loc, loc)", + "evidence": "Obfusc: L120: SETUP_SCRIPT_STUB = \"__import__('setuptools').setup()\"\nExec: L447: exec(finder, loc, loc)", "evidence_hash": "a78d7f5af7eb4ba92656cda258c195b92f6337c585c97d0823e47a9d4a2eb15d" }, { @@ -1322,12 +1394,20 @@ "evidence": "Obfusc: L919: c = compile(funcstr, filename, 'exec')\nExec: L163: module = eval(import_command) | L170: exec(import_command, {}, namespace) | L903: exec(ln, {}, namespace) | L909: exec(ln, {}, namespace) | L920: exec(c, namespace, funclocals)", "evidence_hash": "ab4f5819576a70038301668b8f3e4a781c4b757b146117d5d93eab1896a5a6cd" }, + { + "package": "tensorboard", + "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", + "check": "Python wheel ships large JS bundle (uncommon; manually review)", + "severity": "HIGH", + "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", + "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" + }, { "package": "torch", "file": "torch/_dynamo/bytecode_debugger.py", "check": "Anti-analysis/sandbox evasion + suspicious behavior", "severity": "HIGH", - "evidence": "Anti: L1048: self._old_trace = sys.gettrace() | L1049: sys.settrace(self._settrace_callback) | L1106: sys.settrace(self._old_trace)\nExec: L683: result = eval(arg, frame_globals, eval_locals) | L708: result = eval(cmd, frame_globals, eval_locals) | L716: exec(cmd, frame_globals, eval_locals)", + "evidence": "Anti: L1052: self._old_trace = sys.gettrace() | L1053: sys.settrace(self._settrace_callback) | L1113: sys.settrace(self._old_trace)\nExec: L684: result = eval(arg, frame_globals, eval_locals) | L709: result = eval(cmd, frame_globals, eval_locals) | L717: exec(cmd, frame_globals, eval_locals)", "evidence_hash": "dc2afd1769d357c15b69802bd2799fafa059c0b1dcdd4937528fb5b601962f1b" }, { @@ -1343,7 +1423,7 @@ "file": "torch/fx/experimental/rewriter.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L46: code = compile(dest_ast, \"\", \"exec\")\nExec: L49: exec(code, globals_dict)", + "evidence": "Obfusc: L44: code = compile(dest_ast, \"\", \"exec\")\nExec: L47: exec(code, globals_dict)", "evidence_hash": "76374f96feed416eec390458843621f33524cfb8d93ef0f3eb4cb1b47d0ad748" }, { @@ -1359,7 +1439,7 @@ "file": "torch/package/package_importer.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L602: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", + "evidence": "Obfusc: L599: def __import__(self, name, globals=None, locals=None, fromlist=(), level=0):\nExec: L412: exec(code, ns)", "evidence_hash": "c7c0650f0c74a086d224112f77ee76634b8f47afc047ce27fee8c7fc45560512" }, { @@ -1391,7 +1471,7 @@ "file": "tests/test_mlx_trainer_internals.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L430: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L408: def eval(self):", + "evidence": "Obfusc: L1158: assert ppl == pytest.approx(__import__(\"math\").exp(2.5))\nExec: L1136: def eval(self):", "evidence_hash": "c409327ef6420cc0c7224506fcb82b11bbc9838a6f2f97c9c2cfc00a40c4cdbf" }, { @@ -1407,7 +1487,7 @@ "file": "unsloth_zoo/compiler.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4294: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4291: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4292: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4293: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3385: exec(f\"import {model_location}\", globals()) | L3388: modeling_file = eval(model_location) | L3401: exec(\nL3402: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3403: ) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3407: globals(),\nL3408: locals(),\nL3409: ) | L3560: source = eval(f\"modeling_file.{module}\") | L3574: source = eval(f\"modeling_file.{module}\") | L3675: source = eval(f\"modeling_file.{module}\") | L3713: source = eval(f\"{model_location}.{module}\") | L3784: source = eval(f\"{model_location}.{module}\") | L3832: source = eval(f\"{model_location}.{module}\") | L4054: source = eval(f\"{model_location}.{module}\") | L4065: exec(\nL4066: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4067: globals(),\nL4068: ) | L4131: source = eval(f\"{model_location}.{module}\") | L4172: module_cls = eval(f\"{model_location}.{module}\") | L4209: module_cls = eval(f\"{model_location}.{module}\") | L4276: exec(\nL4277: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4278: globals(),\nL4279: ) | L4341: exec(inner_training_loop, globals()) | L4349: function = eval(f\"{model_location}.{module}\") | L4427: function = eval(f\"{model_location}.{module}\") | L4562: source = eval(f\"{model_location}.torch\") | L4569: function = eval(f\"source.nn.{module}\") | L4628: exec(\nL4629: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4630: globals(),\nL4631: locals(),\nL4632: ) | L4634: exec(\nL4635: f\"{model_location}.nn.{module}.forward = forward\",\nL4636: globals(),\nL4637: locals(),\nL4638: ) | L4642: exec(\nL4643: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4644: globals(),\nL4645: locals(),\nL4646: ) | L4648: exec(\nL4649: f\"combined_module.nn.{module}.forward = forward\",\nL4650: globals(),\nL4651: locals(),\nL4652: ) | L4669: exec(\nL4670: f\"{model_location}.{module} = combined_module.{module}\",\nL4671: globals(),\nL4672: locals(),\nL4673: ) | L4683: check_dicts = dir(eval(f\"{model_location}\")) | L4685: item = eval(f\"{model_location}.{check}\") | L4695: exec(\nL4696: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4697: globals(),\nL4698: locals(),\nL4699: )", + "evidence": "Obfusc: L1013: _mod = __import__(model_location, fromlist=items) | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nL4298: f' \"-____-\" Trainable parameters = {get_model_param_count(model, trainable_only=True):,} of {get_model_param_count(model):,} ({get_model_param_count(model, trainable_only=True)/get_model_p sha256:9e832c1e7b2815aa44dc42638d821cb9df22550db88d85bd4bfb29c13f984b53 | L4295: f\" {chr(92)}{chr(92)} /| Num examples = {num_examples:,} | Num Epochs = {num_train_epochs:,} | Total steps = {max_steps:,}\\\\n\"\\\\\nL4296: f\"O^O/ {chr(92)}_/ {chr(92)} Batch size per device = {self._train_batch_size:,} | Gradient accumulation steps = {args.gradient_accumulation_steps}\\\\n\"\\\\\nL4297: f\"{chr(92)} / Data Parallel GPUs = {args.world_size} | Total batch size ({self._train_batch_size} x {args.gradient_accumulation_steps} x {args.world_size}) = {total_train_batch_size: sha256:b7ea9cdbe360ad323911014caf12a9d4ec87cb36195a511080e4e484c2fb80d2\nExec: L612: if eval(_dtype) is not None: | L613: dtype = eval(_dtype) | L955: _modeling_file = eval(model_location) | L1255: f = eval(f\"{model_location}.{module}\") | L1563: exec(f\"def raise_{j}(*args, **kwargs): print('{function}')\", globals(), locals()) | L1564: try: exec(f\"EMPTY_LOGITS.{function} = raise_{j}\", globals(), locals()) | L2699: exec(f\"import {parent}\", locals(), globals()) | L2830: dir(eval(parent)), | L2834: exec(f\"{parent}.{child}.forward = forward\", globals(), locals()) | L2908: module = eval(f\"modeling_file.{module}\") | L2935: inner_class = eval(f\"modeling_file.{inner_class}\") | L3065: exec(f\"from timm.layers.norm_act import {norm}\") | L3073: forward = eval(norm).forward | L3079: exec(f\"timm.layers.norm_act.{norm}.forward = forward\") | L3096: exec(f\"from timm.models._efficientnet_blocks import {block}\") | L3104: forward = eval(block).forward | L3110: exec(f\"timm.models._efficientnet_blocks.{block}.forward = forward\") | L3389: exec(f\"import {model_location}\", globals()) | L3392: modeling_file = eval(model_location) | L3405: exec(\nL3406: \"model_logger.addFilter(HideLoggingMessage('`use_cache`'))\", globals(), locals()\nL3407: ) | L3409: exec(\nL3410: \"model_logger.addFilter(HideLoggingMessage('compile_config'))\",\nL3411: globals(),\nL3412: locals(),\nL3413: ) | L3564: source = eval(f\"modeling_file.{module}\") | L3578: source = eval(f\"modeling_file.{module}\") | L3679: source = eval(f\"modeling_file.{module}\") | L3717: source = eval(f\"{model_location}.{module}\") | L3788: source = eval(f\"{model_location}.{module}\") | L3836: source = eval(f\"{model_location}.{module}\") | L4058: source = eval(f\"{model_location}.{module}\") | L4069: exec(\nL4070: f\"{model_location}.{module}._update_causal_mask = no_update_causal_mask\",\nL4071: globals(),\nL4072: ) | L4135: source = eval(f\"{model_location}.{module}\") | L4176: module_cls = eval(f\"{model_location}.{module}\") | L4213: module_cls = eval(f\"{model_location}.{module}\") | L4280: exec(\nL4281: \"from transformers.trainer import (\" + \", \".join(x for x in good_items) + \")\",\nL4282: globals(),\nL4283: ) | L4345: exec(inner_training_loop, globals()) | L4353: function = eval(f\"{model_location}.{module}\") | L4431: function = eval(f\"{model_location}.{module}\") | L4566: source = eval(f\"{model_location}.torch\") | L4573: function = eval(f\"source.nn.{module}\") | L4632: exec(\nL4633: f\"{model_location}.torch.nn.{module}.forward = forward\",\nL4634: globals(),\nL4635: locals(),\nL4636: ) | L4638: exec(\nL4639: f\"{model_location}.nn.{module}.forward = forward\",\nL4640: globals(),\nL4641: locals(),\nL4642: ) | L4646: exec(\nL4647: f\"combined_module.torch.nn.{module}.forward = forward\",\nL4648: globals(),\nL4649: locals(),\nL4650: ) | L4652: exec(\nL4653: f\"combined_module.nn.{module}.forward = forward\",\nL4654: globals(),\nL4655: locals(),\nL4656: ) | L4673: exec(\nL4674: f\"{model_location}.{module} = combined_module.{module}\",\nL4675: globals(),\nL4676: locals(),\nL4677: ) | L4687: check_dicts = dir(eval(f\"{model_location}\")) | L4689: item = eval(f\"{model_location}.{check}\") | L4699: exec(\nL4700: f\"{model_location}.{check}['{key}'] = combined_module.{replaced_class}\",\nL4701: globals(),\nL4702: locals(),\nL4703: )", "evidence_hash": "ec1875fd32d00fe885e566ebda75163e46e838ca31020abb57e0991892c2bdf7" }, { @@ -1423,8 +1503,8 @@ "file": "unsloth_zoo/mlx/loader.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L2218: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L140: mx.eval(model.parameters()) | L176: mx.eval(model.parameters()) | L2022: model.eval() | L2605: mx.eval(model.parameters()) | L2721: mx.eval(module.weight) | L4030: mx.eval(model.parameters()) | L4058: mx.eval(model.parameters()) | L4178: mx.eval(model.parameters())", - "evidence_hash": "9b29dade82912216c8b4808aa293b79749aa80ef1d2be35edd93bec7632810f1" + "evidence": "Obfusc: L2869: _mod = __import__(module_name, fromlist=[\"_\"])\nExec: L148: mx.eval(model.parameters()) | L180: mx.eval(model.parameters()) | L732: mx.eval(model.parameters()) | L733: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L799: mx.eval(model.parameters()) | L802: mx.eval(mx.distributed.all_sum(mx.array(1.0), stream=mx.cpu)) | L2673: model.eval() | L3256: mx.eval(model.parameters()) | L3372: mx.eval(module.weight) | L5666: mx.eval(model.parameters()) | L5716: mx.eval(model.parameters()) | L5859: mx.eval(model.parameters())", + "evidence_hash": "7b44760032c5df6d379ccfdd0bff3d23f857f64e08210fa0fba8d2881d457634" }, { "package": "unsloth-zoo", @@ -1439,7 +1519,7 @@ "file": "unsloth_zoo/saving_utils.py", "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", "severity": "HIGH", - "evidence": "Obfusc: L3241: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3123: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3169: exec(save_pretrained, globals(), functions)", + "evidence": "Obfusc: L4015: module = __import__('transformers', fromlist=[model_class_name])\nExec: L3897: exec(f\"from transformers.modeling_utils import ({', '.join(functions)})\", locals(), globals()) | L3943: exec(save_pretrained, globals(), functions)", "evidence_hash": "530b2383acd9fe8330aa65cd0bf86164aaacd47770e7c8d0752195bee36396ec" }, { @@ -1449,118 +1529,6 @@ "severity": "HIGH", "evidence": "Obfusc: L836: code = compile(module, \"\", \"exec\")\nExec: L736: exec(code, globs, locs)", "evidence_hash": "5c0992c90f05c772abd94d00784f157de337e1f8567f8b3aee1b15e46c96cd5d" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3355: os.dup2(conn.fileno(), i) | L3387: \"test needs os.dup2()\") | L3405: os.dup2(fd, newfd) | L19: import socket sha256:26a745abdc7e89da28ab943394234d8ccb415e805477c3cc1f7d4766341a4c4c", - "evidence_hash": "a6b9bb85e9bb6682ab0dea4f95fd9266e8802f118c76d86dd87f7ab5864872cf" - }, - { - "package": "unsloth-zoo", - "file": "scripts/scan_packages.py", - "check": "Writes to /tmp and executes (staged dropper)", - "severity": "CRITICAL", - "evidence": "L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", | L308: r\"/tmp/\\S+.*(?:subprocess|os\\.system|os\\.popen|Popen|chmod.*\\+x)\", sha256:78268349021e21bedcd2eaaa5b4a71b0de1d52e023ada914dfdc09515ee1aad8", - "evidence_hash": "590fe1c96c442fbea5eb8642650257bc0b0199e919b9bacdb11dfa767b6fe839" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3521: os.dup2(conn.fileno(), i) | L3553: \"test needs os.dup2()\") | L3571: os.dup2(fd, newfd) | L20: import socket sha256:07d2933301c0dbeeb6e42381687827d8dd7cfd7471986c559ca64283d5ae6e24", - "evidence_hash": "db1f4ca69865ec3911d7450fe11d212b817139deda21cd7a4ee32d547a8dc452" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L586: while True: sha256:bef9ea429314fad39e063895a37dc5cfe9b04561f3d1acbb3c99abb4e92e6cfe", - "evidence_hash": "b15773e1bc249713156a349278ea60f7c0e3dd7d537affe929ab51089e1942bb" - }, - { - "package": "tensorboard", - "file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js", - "check": "Python wheel ships large JS bundle (uncommon; manually review)", - "severity": "HIGH", - "evidence": "sha256: 53c38430766be25dc672a30846ac3b9eba86aee35eb0746785ec012647c7d9a2", - "evidence_hash": "2c6384e8115a6d5dacf1f84d8f724832d8dc59feb442bb98ffae0857c0ccb381" - }, - { - "package": "fastapi", - "file": "fastapi/routing.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L586: while True: sha256:251135b5ebfdd1248916449f32262575e003ef64382501c65b7e4061d67bda45", - "evidence_hash": "365aef4449c8089753d9398417cd76ab762cef547d75db70d87bca9c0b550ab5" - }, - { - "package": "fastmcp-slim", - "file": "fastmcp/cli/apps_dev.py", - "check": "Enumerates filesystem AND makes network calls", - "severity": "CRITICAL", - "evidence": "FS: L637: history.replaceState(null, \"\", url); sha256:17068ba5bfed62c3a3007ec8bf3e0ea41ef6529b9e6112064d9afb3be9231436\nNetwork: L1304: with httpx.Client(timeout=30.0) as client: | L1318: with httpx.Client(timeout=30.0) as client: | L1348: with httpx.Client(timeout=30.0) as client: | L1549: client = httpx.AsyncClient(\nL1550: timeout=httpx.Timeout(60.0, read=None), trust_env=False\nL1551: ) | L1713: async with httpx.AsyncClient(trust_env=False) as client: | L1781: with socket.socket(family, socket.SOCK_STREAM) as s:", - "evidence_hash": "e5325edfada6499540e6f0c24a0868979d275522e2b6a180aa9b5dd3280681b4" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/_sandbox.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L1179: while True: sha256:33ceddf9e42aae207e891e97808c518e92a0b27ab60e4326256717bfb25a3a38", - "evidence_hash": "802fd41d8bb17bf425e99d128c0351c820103a5efb74690a4086e542a71437b8" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/_sandbox.py", - "check": "Writes to /tmp and executes (staged dropper)", - "severity": "CRITICAL", - "evidence": "L83: d=/tmp/.sbx-server\nL84: if command -v wget >/dev/null 2>&1; then wget -q --header \"Authorization: Bearer $SBX_DL_TOKEN\" -O \"$d\" \"$SBX_SERVER_URL\"\nL85: elif command -v curl >/dev/null 2>&1; then curl -fsSL -H \"Authorization: Bearer $SBX_DL_TOKEN\" -o \"$d\" \"$SBX_SERVER_URL\"\nL86: else cp \"$SBX_SERVER_MOUNT/sbx-server\" \"$d\"; fi\nL87: chmod +x \"$d\"", - "evidence_hash": "6908a3fe328fa94ee22a119998d6ad07cfa1ba4efa2628acf240f4204fd76e22" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/hf_api.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L4613: while True: sha256:f764b6ca3118b23c7c0e670e77178c022a6905f825d7df6e528545fa10aae8f6", - "evidence_hash": "9c85d50c227285fa8dc69512999cbb082258cda4b299c7d0e0f69f5aff7accd4" - }, - { - "package": "huggingface-hub", - "file": "huggingface_hub/utils/_http.py", - "check": "C2 polling/beaconing loop detected", - "severity": "CRITICAL", - "evidence": "L462: while True: sha256:c75d1ee228cf7703a8c28551d649395a1f89f69a3aba69413f5bbcbd10c31958", - "evidence_hash": "d4d5f83fed39b87898cf776d5dad0bf1a6388a932f5fb7997d1070b50e46213e" - }, - { - "package": "cffi", - "file": "cffi/_cffi_gen_src.py", - "check": "Advanced obfuscation (marshal/compile/zlib) + exec/eval", - "severity": "HIGH", - "evidence": "Obfusc: L52: compiled = compile(source=pysrc, filename=filename, mode='exec')\nExec: L53: exec(compiled, globs, globs)", - "evidence_hash": "c429e4c977a61db6b7c717b5a552fce74eda622213e49eb5467a3782fd746fb9" - }, - { - "package": "multiprocess", - "file": "multiprocess/forkserver.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L6: import socket sha256:6c707119169286c9a798e2c8d13a48614e481d8a503950916fd4ffb4c94d3182", - "evidence_hash": "50fec0f0522a8e4e636bf348b752002d7935d8455af31fb78c6f11e2eba19f6d" - }, - { - "package": "multiprocess", - "file": "multiprocess/tests/__init__.py", - "check": "Reverse shell / bind shell pattern", - "severity": "CRITICAL", - "evidence": "L3569: os.dup2(conn.fileno(), i) | L3601: \"test needs os.dup2()\") | L3619: os.dup2(fd, newfd) | L20: import socket sha256:c824dc0f409f242420c3fbb324790c53cb3078d2c8b07ee8f2a05694b01c2946", - "evidence_hash": "3878a2b430c175dbc5877a95195bfe52f9588ff73fb74e2261ed5e33087915ad" } ] } From fb5dc91bb4f33f8a4c5a41c89cbe88c93394e970 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 9 Jul 2026 05:09:16 -0700 Subject: [PATCH 052/402] Studio: remove dead direct_linux_release_plan path (#7030) parse_direct_linux_release_bundle and direct_linux_release_plan are no longer reached by any live code path. Fork Linux installs resolve through _fork_manifest_release_plans -> _linux_published_attempts, and the upstream (ggml-org) path uses direct_upstream_release_plan. The dead parser also called _resolve_linux_bundle_profile, which no longer exists, so its CUDA branch would raise NameError if ever executed. Drop both functions and the obsolete TestDirectLinuxNvidiaCpuGate; its live equivalent TestLinuxPublishedAttemptsNvidiaCpuGate already covers the NVIDIA no-silent-CPU behaviour. --- studio/install_llama_prebuilt.py | 156 ------------------- tests/studio/install/test_selection_logic.py | 62 -------- 2 files changed, 218 deletions(-) diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 856ba71478..40caebc040 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -1286,162 +1286,6 @@ def synthetic_checksums_for_release( ) -def parse_direct_linux_release_bundle( - repo: str, release: dict[str, Any] -) -> PublishedReleaseBundle | None: - release_tag = release.get("tag_name") - if not isinstance(release_tag, str) or not release_tag: - return None - - assets = release_asset_map(release) - artifacts: list[PublishedLlamaArtifact] = [] - inferred_labels: list[str] = [] - - linux_asset_re = re.compile( - r"^app-(?P