From 95f42bcceed7bbecb244adfbc16484cd23a9f13f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 24 Jul 2026 22:34:48 -0700 Subject: [PATCH 01/25] tests: restore the inheritance-before-guard ordering assertion (#7251) The gguf order fix that landed on main dropped the only assertion covering the prerequisite that llama_extra_args inheritance runs before the GGUF branch: the inherited value (a carried --no-mmproj) shapes the hub guard's require_mmproj, so a future reorder could reject a load over an mmproj download the inherited arguments would disable. The comment also misattributed the inheritance site to _guard_chat_load_against_training. The assertion is restored anchored on the call form "= _resolve_inherited_extra_args(", which pins the endpoint's call site (the bare name would match the function definition, which always precedes the endpoint, making the check vacuous), and the comment now names the real inheritance site. 32 tests pass. --- studio/backend/tests/test_gguf_load_cache_reuse.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index 6d1fac980b..dfd6fc6034 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -789,9 +789,12 @@ class TestLoadHubDownloadExclusion: # The gguf_load_in_flight marker must be entered before the hub-download # guard and the unload so a concurrent load can't race the download - # manager. The llama_extra_args inheritance that used to sit between the - # marker and the guard now runs in _guard_chat_load_against_training, ahead - # of the GGUF branch, so it is no longer a landmark inside this slice. + # manager. The llama_extra_args inheritance moved out of the branch into + # _resolve_inherited_extra_args, which must still run BEFORE it: the + # inherited value (e.g. a carried --no-mmproj) shapes the guard's + # require_mmproj. Anchor on the call form so the assertion pins the + # endpoint's call site, not the function definition. + assert source.index("= _resolve_inherited_extra_args(") < source.index("if config.is_gguf:") assert ( gguf_branch.index("enter_context(gguf_load_in_flight") < gguf_branch.index("_hub_download_blocks_gguf_load") From dffea2af27b9ee5a10479fb6566184394d6921cd Mon Sep 17 00:00:00 2001 From: Souravrajvi0 <144546710+Souravrajvi0@users.noreply.github.com> Date: Sat, 25 Jul 2026 11:57:47 +0530 Subject: [PATCH 02/25] fix(studio): honor run settings on initial model load (#7346) (#7351) * fix(studio): honor run settings on initial model load When loading a model from the gear-icon run-settings page, Context Length and KV Cache Dtype were ignored if the user clicked Load before blurring the context field, or before React flushed staged config into the store. - Add NumericValueInput.commit() to flush a focused draft on Load - Pass effectiveLoadConfig from model-config-page to onRun - Prefer selection.config in performLoad for all load knobs - Preserve meta.forceReload from the config-page reload path Fixes #7346 * fix(studio): flush NumericValueInput draft when Load blurs first Clicking Load blurs the context field before handleRun runs, so commit() returned the stale value prop. Keep draft in a ref and parse it even when the input is no longer focused. * fix(studio): preserve Auto context when Load is clicked without edits NumericValueInput.commit() now returns null unless the user actually changed the field, so GGUF Load/Save no longer pins the displayed native context into customContextLength when Auto was left untouched. * fix(studio): clear NumericValueInput dirty state after blur commit After a normal blur commit, reset dirtyRef so a later Load cannot replay a stale draftRef when the user changed context via Reset or the slider. * test(studio): pin NumericValueInput Auto/dirty contracts for #7346 Lock Codex P1/P2: commit returns null unless dirty, blur clears dirtyRef, and handleRun only promotes a non-null committed context. * fix(studio): keep same-click context draft after blur (#7346) Blur can commit and clear dirtyRef before Load's onClick; stash that committed value for one imperative commit() so typed context is not lost. * chore: refresh PR head for #7351 * fix(studio): handle context commit edge cases * chore: refresh PR head * test(studio): guard invalid context drafts * style(studio): format context draft guard * test(studio): exercise same-click model config loads * fix(studio): drop stale blur pin when the typed context equals the shown value NumericValueInput cached every blur commit in lastBlurCommittedRef, even when the draft equalled the current value and no onChange was dispatched. Because the displayed value never changed, the useEffect([value]) clear never fired, so a later Reset or external edit that leaves the shown value unchanged could not drop the cache and the next commit() replayed it into an override that Reset had removed. Only cache the blur result when it actually dispatched onChange (final !== value); when final === value the parent is already current and there is nothing to bridge. Add a Playwright regression that re-types the shown context and asserts no override is stored. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: commit every same-click numeric draft before staging the load config The run-settings Load/Reload button flushed only the GGUF Context Length draft imperatively before building the load config. Max Seq Length (non-GGUF), GPU Layers and MoE Layers on CPU (GGUF) are the same NumericValueInput and stage their typed value only on blur, so editing one and clicking Load in the same gesture staged the load from a still-stale parent config and dropped the value the user just typed. Wire an imperative commit handle through those inputs too and fold every committed draft into the effective config, recomputing the non-GGUF load-time max sequence length from the committed draft. * fix(studio): recompute fixed-layer context pin and drop stale blur cache on every render Two run-settings edge cases on the model-config page: 1) pinFixedLayerContext was computed from the render-time config, before a same-click GPU Layers draft is committed in handleRun. Typing a positive fixed-layer value on an auto-fit GGUF and clicking Reload therefore built the runtime config with customContextLength: null, so a later fresh load sent the native context with fixed layers (the OOM the pin exists to avoid). Recompute the pin from the committed effectiveConfig. 2) NumericValueInput cleared its blur bridge only on a value change. A real edit (final !== value) that Reset then reverts to the same shown number nets value back unchanged, so the effect never re-ran and the stale pin survived into the next Load/Save, replaying the override Reset removed. The bridge is only valid across the single synchronous same-click gesture that set it, so clear it on every settled render instead. Add source-contract regressions for both. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: Unsloth Co-authored-by: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Co-authored-by: Daniel Han Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../frontend/src/features/chat/chat-page.tsx | 2 +- .../chat/hooks/use-chat-model-runtime.ts | 78 +++++--- .../components/model-config-page.tsx | 128 ++++++++++++-- .../components/numeric-value-input.tsx | 166 ++++++++++++++---- .../src/features/model-picker/index.ts | 1 + tests/studio/playwright_model_config.py | 69 +++++++- tests/studio/test_model_picker_contracts.py | 117 +++++++++++- 7 files changed, 479 insertions(+), 82 deletions(-) diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 7452cf3447..a439a91239 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -2676,7 +2676,7 @@ export function ChatPage({ config: meta?.config, nativePathToken: meta?.nativePathToken, nativePathExpiresAtMs: meta?.nativePathExpiresAtMs, - forceReload: isSameLoadedModel || undefined, + forceReload: meta?.forceReload ?? (isSameLoadedModel || undefined), }; await stageOrLoad(selection); })(); 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 bb19223a6a..48a6168555 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 @@ -62,6 +62,7 @@ import { import { isExternalModelId } from "../external-providers"; import { applyPerModelConfigToRuntime, + normalizeMaxSeqLength, type PerModelConfig, } from "@/features/model-picker"; import type { @@ -604,12 +605,19 @@ export function useChatModelRuntime() { async function performLoad(): Promise { if (abortCtrl.signal.aborted) throw new Error("Cancelled"); let previousWasUnloaded = false; + const pendingLoadConfig = + typeof selection !== "string" ? selection.config : undefined; + if (pendingLoadConfig) { + applyPerModelConfigToRuntime(pendingLoadConfig); + } const currentCheckpoint = useChatRuntimeStore.getState().params.checkpoint; const stateBeforeUnload = useChatRuntimeStore.getState(); let trustRemoteCode = stateBeforeUnload.params.trustRemoteCode ?? false; let approvedRemoteCodeFingerprint: string | null = null; - const maxSeqLength = stateBeforeUnload.params.maxSeqLength; + const maxSeqLength = + normalizeMaxSeqLength(pendingLoadConfig?.maxSeqLength) ?? + stateBeforeUnload.params.maxSeqLength; const previousActiveNativePathToken = stateBeforeUnload.activeNativePathToken; const previousIsGguf = @@ -643,34 +651,54 @@ export function useChatModelRuntime() { const previousActiveNativePathExpiresAtMs = stateBeforeUnload.activeNativePathExpiresAtMs; // Snapshot the load settings at click time, before the awaits below - // (validation, the trust dialog, unload). - const loadChatTemplateOverride = stateBeforeUnload.chatTemplateOverride; - const loadKvCacheDtype = stateBeforeUnload.kvCacheDtype; + // (validation, the trust dialog, unload). When the picker staged a + // config payload, prefer it over the store: React may not have + // flushed NumericValueInput's blur commit into state yet. + const loadChatTemplateOverride = + pendingLoadConfig?.chatTemplateOverride?.trim() + ? pendingLoadConfig.chatTemplateOverride + : stateBeforeUnload.chatTemplateOverride; + const loadKvCacheDtype = + pendingLoadConfig?.kvCacheDtype ?? stateBeforeUnload.kvCacheDtype; // gpuMemoryMode is a standing preference (kept across a model switch); // the rest are per-model knobs the reset below clears, so they are // re-baselined there in lock-step with the store. - let loadCustomContextLength = stateBeforeUnload.customContextLength; + let loadCustomContextLength = + pendingLoadConfig?.customContextLength ?? + stateBeforeUnload.customContextLength; const loadGgufContextLength = stateBeforeUnload.ggufContextLength; - const loadTensorParallel = stateBeforeUnload.tensorParallel; + const loadTensorParallel = + pendingLoadConfig?.tensorParallel ?? stateBeforeUnload.tensorParallel; const loadActivePresetSource = stateBeforeUnload.activePresetSource; const loadActiveGgufVariant = stateBeforeUnload.activeGgufVariant; - const loadGpuMemoryMode = stateBeforeUnload.gpuMemoryMode; - let loadGpuLayers = stateBeforeUnload.gpuLayers; - let loadNCpuMoe = stateBeforeUnload.nCpuMoe; + const loadGpuMemoryMode = + pendingLoadConfig?.gpuMemoryMode ?? stateBeforeUnload.gpuMemoryMode; + let loadGpuLayers = + pendingLoadConfig?.gpuLayers ?? stateBeforeUnload.gpuLayers; + let loadNCpuMoe = + pendingLoadConfig?.nCpuMoe ?? stateBeforeUnload.nCpuMoe; let loadSplitRatio = stateBeforeUnload.splitRatio; // Reconcile the persisted pick against the GPUs present now, so a stale // cross-host / now-hidden pick is dropped before /load rather than // rejected there. Warm the device cache first: load-on-selection can // run before any GPU hook mounted, and a cold cache would pass the // pick through unvalidated. validateGpuIds derives from this too. - if (stateBeforeUnload.selectedGpuIds != null) { + if ( + pendingLoadConfig?.selectedGpuIds !== undefined || + stateBeforeUnload.selectedGpuIds != null + ) { await ensureGpuDeviceCache(); } - let loadSelectedGpuIds = reconcilePersistedGpuIds( - stateBeforeUnload.selectedGpuIds, - ); - let loadSpeculativeType = stateBeforeUnload.speculativeType; - let loadSpecDraftNMax = stateBeforeUnload.specDraftNMax; + let loadSelectedGpuIds = + pendingLoadConfig?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds) + : reconcilePersistedGpuIds(stateBeforeUnload.selectedGpuIds); + let loadSpeculativeType = + pendingLoadConfig?.speculativeType != null + ? normalizeSpeculativeType(pendingLoadConfig.speculativeType) + : stateBeforeUnload.speculativeType; + let loadSpecDraftNMax = + pendingLoadConfig?.specDraftNMax ?? stateBeforeUnload.specDraftNMax; try { // Lightweight pre-flight validation: avoid unloading a working model // if the new identifier is clearly invalid (e.g. bad HF id / path). @@ -810,15 +838,23 @@ export function useChatModelRuntime() { // model loads at Auto/native, not the previous model's pin. customContextLength: null, }); - loadSpeculativeType = persistedSpeculativeType; - loadSpecDraftNMax = null; + loadSpeculativeType = + pendingLoadConfig?.speculativeType != null + ? normalizeSpeculativeType(pendingLoadConfig.speculativeType) + : persistedSpeculativeType; + loadSpecDraftNMax = pendingLoadConfig?.specDraftNMax ?? null; // Keep the click-time snapshot in lock-step with the store reset so // the load below sizes against the cleared per-model knobs, not the // previous model's (gpuMemoryMode is standing, so left as captured). - loadCustomContextLength = null; - loadSelectedGpuIds = null; - loadGpuLayers = GPU_LAYERS_AUTO; - loadNCpuMoe = 0; + // An explicit staged config from run-settings still wins. + loadCustomContextLength = + pendingLoadConfig?.customContextLength ?? null; + loadSelectedGpuIds = + pendingLoadConfig?.selectedGpuIds !== undefined + ? reconcilePersistedGpuIds(pendingLoadConfig.selectedGpuIds) + : null; + loadGpuLayers = pendingLoadConfig?.gpuLayers ?? GPU_LAYERS_AUTO; + loadNCpuMoe = pendingLoadConfig?.nCpuMoe ?? 0; loadSplitRatio = null; } diff --git a/studio/frontend/src/features/model-picker/components/model-config-page.tsx b/studio/frontend/src/features/model-picker/components/model-config-page.tsx index afbd33af7c..90202a2bcf 100644 --- a/studio/frontend/src/features/model-picker/components/model-config-page.tsx +++ b/studio/frontend/src/features/model-picker/components/model-config-page.tsx @@ -24,7 +24,14 @@ import { ChevronDownStandardIcon } from "@/lib/chevron-icons"; import { toast } from "@/lib/toast"; import { ArrowLeft01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { type ReactNode, useEffect, useId, useState } from "react"; +import { + type ReactNode, + type Ref, + useEffect, + useId, + useRef, + useState, +} from "react"; import { useDefaultChatTemplate, useModelMaxPositionEmbeddings, @@ -50,7 +57,10 @@ import { } from "../model-config/per-model-config"; import { ChatTemplateEditorDialog } from "./chat-template-editor-dialog"; import type { ModelPickTarget } from "./model-selector/types"; -import { NumericValueInput } from "./numeric-value-input"; +import { + NumericValueInput, + type NumericValueInputHandle, +} from "./numeric-value-input"; const ROW_CLASS = "flex min-h-8 items-center justify-between gap-3"; const LABEL_CLASS = @@ -130,11 +140,13 @@ function MaxSeqLengthSetting({ max, inputMax, onChange, + inputRef, }: { value: number; max: number; inputMax: number; onChange: (value: number) => void; + inputRef?: Ref; }) { return (
@@ -146,6 +158,7 @@ function MaxSeqLengthSetting({
void; displayValue?: string; info?: ReactNode; + inputRef?: Ref; }) { return (
@@ -199,6 +214,7 @@ function AdvancedGpuSlider({ {info && {info}}
) => void; layerCount: number | null; moeLayerCount: number | null; + gpuLayersInputRef?: Ref; + moeLayersInputRef?: Ref; }) { const gpuDevices = useGpuDevices(); const mode = config.gpuMemoryMode ?? "auto"; @@ -322,6 +342,7 @@ function GpuMemorySettings({ <> ) => void; @@ -407,6 +431,8 @@ function GgufAdvancedSettings({ onEditTemplate: () => void; layerCount: number | null; moeLayerCount: number | null; + gpuLayersInputRef?: Ref; + moeLayersInputRef?: Ref; }) { return ( <> @@ -535,6 +561,8 @@ function GgufAdvancedSettings({ update={update} layerCount={layerCount} moeLayerCount={moeLayerCount} + gpuLayersInputRef={gpuLayersInputRef} + moeLayersInputRef={moeLayersInputRef} /> @@ -597,6 +625,10 @@ export function ModelConfigPage({ const [showAdvanced, setShowAdvanced] = useState(() => hasNonDefaultAdvanced(config), ); + const contextInputRef = useRef(null); + const maxSeqLengthInputRef = useRef(null); + const gpuLayersInputRef = useRef(null); + const moeLayersInputRef = useRef(null); const nativePathToken = target.meta.nativePathToken ?? (isActiveModel ? activeNativePathToken : null); @@ -744,11 +776,6 @@ export function ModelConfigPage({ ? { ...config, customContextLength: activeLoadedContext } : config : config; - // Load request needs a concrete max length; substitute the fallback here only, - // never in the persisted runtimeConfig. - const loadConfig = target.isGguf - ? runtimeConfig - : { ...runtimeConfig, maxSeqLength: maxSeqLengthValue }; const rememberChanged = remember !== savedRemember; const persistenceOnly = isActiveModel && atBaseline && rememberChanged; const primaryActionLabel = persistenceOnly @@ -760,18 +787,90 @@ export function ModelConfigPage({ : "Load model"; const handleRun = () => { - const defaultConfig = isDefaultConfig(runtimeConfig); + // Same-click Load/Reload: a numeric draft the user just typed is flushed only + // by that input's blur handler, which updates the parent config after this + // click closure already captured the stale value. Commit every numeric input + // imperatively so the staged load honors what the user just typed, not just + // the Context field. + const committedContext = target.isGguf + ? contextInputRef.current?.commit() + : undefined; + const committedMaxSeqLength = target.isGguf + ? undefined + : maxSeqLengthInputRef.current?.commit(); + const committedGpuLayers = target.isGguf + ? gpuLayersInputRef.current?.commit() + : undefined; + const committedMoeLayers = target.isGguf + ? moeLayersInputRef.current?.commit() + : undefined; + + const pendingPatch: Partial = {}; + if (committedContext != null) { + pendingPatch.customContextLength = committedContext; + } + if (committedMaxSeqLength != null) { + pendingPatch.maxSeqLength = clampMaxSeqLength( + committedMaxSeqLength, + MAX_SEQ_LENGTH_MAX, + ); + } + if (committedGpuLayers != null) { + pendingPatch.gpuLayers = committedGpuLayers; + } + if (committedMoeLayers != null) { + pendingPatch.nCpuMoe = committedMoeLayers; + } + const hasPending = + committedContext != null || + committedMaxSeqLength != null || + committedGpuLayers != null || + committedMoeLayers != null; + + const effectiveConfig = hasPending + ? { ...config, ...pendingPatch } + : config; + // pinFixedLayerContext above was computed from the render-time config, before + // the same-click GPU Layers draft was committed. Recompute it from + // effectiveConfig so committing a positive fixed-layer value still pins the + // fitted context; otherwise the saved config carries customContextLength: null + // and a later fresh load sends the native context with fixed layers (the OOM + // the pin exists to avoid). + const effectivePinFixedLayerContext = + target.isGguf && + effectiveConfig.gpuMemoryMode === "manual" && + effectiveConfig.gpuLayers != null && + effectiveConfig.gpuLayers >= 0 && + effectiveConfig.customContextLength == null && + activeLoadedContext != null; + const effectiveRuntimeConfig = hasPending + ? effectivePinFixedLayerContext + ? { ...effectiveConfig, customContextLength: activeLoadedContext } + : effectiveConfig + : runtimeConfig; + // Non-GGUF load substitutes the resolved max sequence length; recompute it + // from the committed draft so a same-click Max Seq Length edit is not lost. + const effectiveMaxSeqLengthValue = + committedMaxSeqLength == null + ? maxSeqLengthValue + : (normalizeMaxSeqLength(effectiveConfig.maxSeqLength) ?? + clampMaxSeqLength(DEFAULT_MAX_SEQ_LENGTH, nativeMaxSeqLength)); + // Recheck the committed draft so Save/Forget reloads when needed. + const effectiveAtBaseline = perModelConfigsEqual(effectiveConfig, baseline); + const effectivePersistenceOnly = + isActiveModel && effectiveAtBaseline && rememberChanged; + const defaultConfig = isDefaultConfig(effectiveRuntimeConfig); let saveFailed = false; if (remember) { saveFailed = !savePerModelConfig( target.id, target.ggufVariant, - runtimeConfig, + effectiveRuntimeConfig, ); } else { saveFailed = !deletePerModelConfig(target.id, target.ggufVariant); } - if (persistenceOnly) { + if (effectivePersistenceOnly) { if (saveFailed) { toast.error("Couldn't save settings for this model."); return; @@ -791,7 +890,10 @@ export function ModelConfigPage({ if (saveFailed) { toast.error("Couldn't save these settings, loading with them anyway."); } - onRun(loadConfig); + const effectiveLoadConfig = target.isGguf + ? effectiveRuntimeConfig + : { ...effectiveRuntimeConfig, maxSeqLength: effectiveMaxSeqLengthValue }; + onRun(effectiveLoadConfig); }; return ( @@ -838,6 +940,7 @@ export function ModelConfigPage({ setTemplateOpen(true)} layerCount={stagedDims?.layerCount ?? null} moeLayerCount={stagedDims?.moeLayerCount ?? null} + gpuLayersInputRef={gpuLayersInputRef} + moeLayersInputRef={moeLayersInputRef} /> )} @@ -914,6 +1019,7 @@ export function ModelConfigPage({ value={maxSeqLengthValue} max={maxSeqLengthMax} inputMax={MAX_SEQ_LENGTH_MAX} + inputRef={maxSeqLengthInputRef} onChange={(value) => update({ maxSeqLength: clampMaxSeqLength(value, MAX_SEQ_LENGTH_MAX), diff --git a/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx b/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx index 2489927fc2..be9aa7745c 100644 --- a/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx +++ b/studio/frontend/src/features/model-picker/components/numeric-value-input.tsx @@ -2,7 +2,13 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { cn } from "@/lib/utils"; -import { useRef, useState } from "react"; +import { + forwardRef, + useEffect, + useImperativeHandle, + useRef, + useState, +} from "react"; export function snapToStep( value: number, @@ -28,44 +34,109 @@ function sanitizeNumeric(raw: string, allowNegative: boolean): string { return `${sign}${head}${tail}`; } -export function NumericValueInput({ - value, - min, - max, - step, - onChange, - displayValue, - className, - ariaLabel, - size: sizeAttr, - disabled = false, -}: { - value: number; - min?: number; - max?: number; - step: number; - onChange: (v: number) => void; - displayValue?: string; - className?: string; - ariaLabel?: string; - size?: number; - disabled?: boolean; -}) { +export type NumericValueInputHandle = { + /** Commit a valid focused/same-click draft; null when none is pending. */ + commit: () => number | null; +}; + +export const NumericValueInput = forwardRef< + NumericValueInputHandle, + { + value: number; + min?: number; + max?: number; + step: number; + onChange: (v: number) => void; + displayValue?: string; + className?: string; + ariaLabel?: string; + size?: number; + disabled?: boolean; + } +>(function NumericValueInput( + { + value, + min, + max, + step, + onChange, + displayValue, + className, + ariaLabel, + size: sizeAttr, + disabled = false, + }, + ref, +) { const [focused, setFocused] = useState(false); const [draft, setDraft] = useState(""); const cancelBlurCommitRef = useRef(false); + const draftRef = useRef(""); + const dirtyRef = useRef(false); + // Same-click Load: blur commits via onChange and clears dirtyRef before the + // button onClick runs, while parent `value` is still stale. Keep the blur + // result for one imperative commit(); clear when `value` catches up or on + // focus / external edits (Reset, slider). + const lastBlurCommittedRef = useRef(null); - const commit = (raw: string) => { + // The blur bridge is only valid across the single synchronous gesture that set + // it: blur commits during a button's mousedown and that button's onClick + // consumes it via commit() before React re-renders. Any settled render means the + // gesture is over, so drop the cache on every commit. Keying this on [value] + // alone missed a Reset (or other external edit) that restores the shown value + // unchanged when the blur did dispatch onChange (final !== value): value nets + // back to its prior number, so the effect never re-ran, the stale pin survived, + // and the next Load/Save replayed the override Reset had removed. + useEffect(() => { + lastBlurCommittedRef.current = null; + }); + + const commitDraft = (raw: string): number | null => { const parsed = Number.parseFloat(raw); if (!Number.isFinite(parsed)) { - return; + return null; } const final = snapToStep(parsed, step, min, max); if (final !== value) { onChange(final); } + return final; }; + useImperativeHandle( + ref, + () => ({ + commit: () => { + if (dirtyRef.current) { + const raw = draftRef.current; + const final = commitDraft(raw); + dirtyRef.current = false; + lastBlurCommittedRef.current = null; + if (final == null) { + draftRef.current = String(value); + } + if (focused) { + setFocused(false); + } + return final; + } + const blurCommitted = lastBlurCommittedRef.current; + if (blurCommitted != null) { + lastBlurCommittedRef.current = null; + if (focused) { + setFocused(false); + } + return blurCommitted; + } + if (focused) { + setFocused(false); + } + return null; + }, + }), + [draft, focused, max, min, onChange, step, value], + ); + const displayed = focused ? draft : (displayValue ?? String(value)); return ( @@ -82,7 +153,11 @@ export function NumericValueInput({ aria-label={ariaLabel} onFocus={(e) => { cancelBlurCommitRef.current = false; - setDraft(String(value)); + dirtyRef.current = false; + lastBlurCommittedRef.current = null; + const next = String(value); + draftRef.current = next; + setDraft(next); setFocused(true); const target = e.currentTarget; requestAnimationFrame(() => target.select()); @@ -90,24 +165,47 @@ export function NumericValueInput({ onBlur={() => { if (cancelBlurCommitRef.current) { cancelBlurCommitRef.current = false; - } else { - commit(draft); + lastBlurCommittedRef.current = null; + } else if (dirtyRef.current) { + const final = commitDraft(draftRef.current); + dirtyRef.current = false; + if (final == null) { + draftRef.current = String(value); + lastBlurCommittedRef.current = null; + } else { + draftRef.current = String(final); + // Only bridge the still-stale parent value when the blur actually + // dispatched onChange (final !== value). When final === value the + // parent is already current, so there is nothing to bridge; caching + // here would leave a stale pin that a later Reset or external edit + // (which doesn't change the displayed value) can never clear, so a + // following Load/Save would recreate the override Reset removed. + lastBlurCommittedRef.current = final !== value ? final : null; + } } setFocused(false); }} - onChange={(e) => - setDraft(sanitizeNumeric(e.target.value, (min ?? 0) < 0)) - } + onChange={(e) => { + dirtyRef.current = true; + lastBlurCommittedRef.current = null; + const next = sanitizeNumeric(e.target.value, (min ?? 0) < 0); + draftRef.current = next; + setDraft(next); + }} onKeyDown={(e) => { if (e.key === "Enter") { e.currentTarget.blur(); } else if (e.key === "Escape") { cancelBlurCommitRef.current = true; - setDraft(String(value)); + dirtyRef.current = false; + lastBlurCommittedRef.current = null; + const next = String(value); + draftRef.current = next; + setDraft(next); e.currentTarget.blur(); } }} className={cn(className)} /> ); -} +}); diff --git a/studio/frontend/src/features/model-picker/index.ts b/studio/frontend/src/features/model-picker/index.ts index d2b4785ec3..383d441f09 100644 --- a/studio/frontend/src/features/model-picker/index.ts +++ b/studio/frontend/src/features/model-picker/index.ts @@ -12,6 +12,7 @@ export { export { hfModelFitsDevice } from "./components/model-selector/recommended-fit"; export { NumericValueInput, + type NumericValueInputHandle, snapToStep, } from "./components/numeric-value-input"; export { SidebarModelConfig } from "./components/sidebar-model-config"; diff --git a/tests/studio/playwright_model_config.py b/tests/studio/playwright_model_config.py index a8d143a253..d72dc37955 100644 --- a/tests/studio/playwright_model_config.py +++ b/tests/studio/playwright_model_config.py @@ -70,6 +70,8 @@ ART_DIR = os.environ.get("PW_ART_DIR", "logs/playwright_modelcfg") ART = Path(ART_DIR) ART.mkdir(parents = True, exist_ok = True) STRICT = os.environ.get("STUDIO_UI_STRICT", "0") == "1" +PLAYWRIGHT_BROWSER = os.environ.get("STUDIO_PLAYWRIGHT_BROWSER", "chromium").lower() +PLAYWRIGHT_CHANNEL = os.environ.get("STUDIO_PLAYWRIGHT_CHANNEL") or None TURN_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_TURN_TIMEOUT_MS", "180000")) WALL_TIMEOUT_S = float(os.environ.get("STUDIO_UI_WALL_TIMEOUT_S", "720")) FETCH_TIMEOUT_MS = int(os.environ.get("STUDIO_UI_FETCH_TIMEOUT_MS", "30000")) @@ -145,10 +147,19 @@ with sync_playwright() as p: ) # Health pre-flight: bash-side health wait can pass before the auth DB migrates. wait_for_health(BASE, timeout = 30.0, info = info) - browser = p.chromium.launch( - headless = True, - args = chromium_launch_args(), - ) + if PLAYWRIGHT_BROWSER not in ("chromium", "firefox", "webkit"): + fail(f"unsupported STUDIO_PLAYWRIGHT_BROWSER={PLAYWRIGHT_BROWSER!r}") + sys.exit(1) + browser_type = getattr(p, PLAYWRIGHT_BROWSER) + launch_kwargs = {"headless": True} + if PLAYWRIGHT_BROWSER == "chromium": + launch_kwargs["args"] = chromium_launch_args() + if PLAYWRIGHT_CHANNEL: + launch_kwargs["channel"] = PLAYWRIGHT_CHANNEL + elif PLAYWRIGHT_CHANNEL: + fail("STUDIO_PLAYWRIGHT_CHANNEL requires chromium") + sys.exit(1) + browser = browser_type.launch(**launch_kwargs) ctx = browser.new_context( viewport = {"width": 1280, "height": 900}, reduced_motion = "reduce", @@ -464,11 +475,6 @@ with sync_playwright() as p: else: default_ctx = ctx_in.input_value() info(f"default Context Length shown: {default_ctx!r}") - ctx_in.click() - ctx_in.fill(str(DISTINCT_CTX)) - page.wait_for_timeout(300) - page.keyboard.press("Tab") # blur to commit - page.wait_for_timeout(300) remember = popover.get_by_label("Remember for this model").first if _count(remember): try: @@ -477,12 +483,16 @@ with sync_playwright() as p: remember.click() else: fail("'Remember for this model' checkbox not found") + ctx_in.click() + ctx_in.fill(str(DISTINCT_CTX)) page.wait_for_timeout(300) shoot("05-ctx-set") btn = primary_button(popover) if btn is None: fail("primary Load/Save button not found in run-settings") else: + # Keep the input focused. The button click must commit the draft + # and use it in the same load request. btn.click() page.wait_for_timeout(2500) shoot("06-after-load") @@ -570,6 +580,47 @@ with sync_playwright() as p: else: info("OK reset: distinctive context cleared from unsloth_model_configs") shoot("08-after-reset") + + # ───────────────────────────────────────────────────── + # 3b. Re-typing the value already shown must not pin an override (HARD). + # Entering the currently displayed native/default context commits no + # onChange (the value is unchanged), so the cached blur value must not be + # replayed into a stored override on Load. Otherwise re-typing the shown + # number, or doing so before a Reset, recreates a phantom context pin. + # ───────────────────────────────────────────────────── + step("re-typing the shown context does not pin an override") + ctx_in = context_input(popover) + native_default = _as_int(ctx_in.input_value()) if ctx_in else None + if ctx_in is None or native_default is None: + info("skip re-type-shown: Context Length input has no numeric default") + else: + remember = popover.get_by_label("Remember for this model").first + if _count(remember): + try: + remember.check() + except Exception: + remember.click() + ctx_in.click() + ctx_in.fill(str(native_default)) + page.wait_for_timeout(200) + btn = primary_button(popover) + if btn is not None and btn.is_enabled(): + # Same-click Load: the button click must commit the draft, but a draft + # equal to the shown value carries no override. + btn.click() + page.wait_for_timeout(1500) + cfg = read_configs() + pinned = any( + _as_int(e.get("customContextLength")) == native_default for e in config_entries(cfg) + ) + if pinned: + fail( + "re-typing the shown context pinned it as an override " + f"(customContextLength={native_default})" + ) + else: + info("OK re-type-shown: shown context not stored as an override") + shoot("08b-after-retype-shown") close_picker() # ───────────────────────────────────────────────────── diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index 3934b72d83..815f68e010 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -346,6 +346,36 @@ def test_fixed_layer_gguf_pins_displayed_context(): assert "customContextLength: activeLoadedContext" in src +def test_fixed_layer_pin_recomputed_after_committing_gpu_layers(): + """pinFixedLayerContext is computed from the render-time config, before a + same-click GPU Layers draft is committed. handleRun must recompute it from the + committed effectiveConfig; otherwise typing a positive GPU Layers value on an + auto-fit GGUF and clicking Reload saves customContextLength: null, so a later + fresh load sends the native context with fixed layers (the OOM the pin avoids).""" + src = _read("features/model-picker/components/model-config-page.tsx") + assert "const effectivePinFixedLayerContext =" in src + assert 'effectiveConfig.gpuMemoryMode === "manual"' in src + assert "effectiveConfig.gpuLayers != null" in src + assert "effectiveConfig.customContextLength == null" in src + assert "{ ...effectiveConfig, customContextLength: activeLoadedContext }" in src + + +def test_blur_cache_cleared_on_every_settled_render(): + """The lastBlurCommittedRef bridge is valid only across the single synchronous + same-click gesture that set it. Keying its clear on [value] missed a Reset (or + external edit) that restores the shown value unchanged after the blur dispatched + onChange: value nets back to its prior number, the effect never re-ran, and a + later Load/Save replayed the override Reset removed. Clear it on every settled + render instead.""" + src = _read("features/model-picker/components/numeric-value-input.tsx") + # The clearing effect must run on every commit, not be gated on [value] alone. + assert not re.search(r"lastBlurCommittedRef\.current = null;\s*\}, \[value\]\);", src) + assert re.search( + r"useEffect\(\(\) => \{\s*lastBlurCommittedRef\.current = null;\s*\}\);", + src, + ) + + def test_auto_defaults_not_persisted_as_overrides(): """Auto GPU memory mode and Auto/default speculative type are follow-global defaults; normalization must not persist them as per-model overrides, else a @@ -382,15 +412,90 @@ def test_reset_persists_null_max_length_and_substitutes_only_for_load(): Reset) so isDefaultConfig can clear a remembered override; the concrete fallback is substituted only into the load request, not the saved record.""" src = _read("features/model-picker/components/model-config-page.tsx") - # Load-only substitution of the resolved value. - assert "maxSeqLength: maxSeqLengthValue" in src - assert "const loadConfig" in src - # The persisted record is loaded via onRun(loadConfig), and save uses the - # untouched runtimeConfig (so a reset/default config stays default). - assert "onRun(loadConfig)" in src + # Load-only substitution of the resolved value (recomputed from any committed + # same-click Max Seq Length draft, so it is never dropped). + assert "maxSeqLength: effectiveMaxSeqLengthValue" in src + assert "const effectiveLoadConfig" in src + # The persisted record is saved from effectiveRuntimeConfig; the load request + # carries effectiveLoadConfig (with any committed context input). + assert "onRun(effectiveLoadConfig)" in src assert "savePerModelConfig(" in src +def test_initial_load_uses_staged_config_payload(): + """Run-settings Load must pass the staged config through to /load even when + React has not flushed NumericValueInput blur commits into the store yet.""" + runtime = _read("features/chat/hooks/use-chat-model-runtime.ts") + assert "const pendingLoadConfig =" in runtime + assert "pendingLoadConfig?.kvCacheDtype" in runtime + assert "pendingLoadConfig?.customContextLength" in runtime + page = _read("features/model-picker/components/model-config-page.tsx") + assert "contextInputRef" in page + assert "contextInputRef.current?.commit()" in page + numeric = _read("features/model-picker/components/numeric-value-input.tsx") + assert "export type NumericValueInputHandle" in numeric + assert "commit:" in numeric + # P1: commit returns null unless the user actually edited the field, + # so Load/Save with untouched Auto does not pin native context. + assert "dirtyRef.current" in numeric + assert "return null;" in numeric + # P2: blur clears dirtyRef after commit so Reset/slider cannot be + # overwritten by a stale draft on a later Load. + assert "dirtyRef.current = false;" in numeric + assert "draftRef.current = String(final);" in numeric + # Same-click Load after blur still sees the committed draft. + assert "lastBlurCommittedRef" in numeric + # Invalid drafts must not turn Auto into an explicit pin. + assert "const commitDraft = (raw: string): number | null" in numeric + assert re.search(r"if \(!Number\.isFinite\(parsed\)\) \{\s*return null;", numeric) + assert re.search( + r"if \(final == null\) \{\s*" + r"draftRef\.current = String\(value\);\s*" + r"lastBlurCommittedRef\.current = null;", + numeric, + ) + # handleRun only promotes commit() when non-null. + assert "committedContext != null" in page + assert "pendingPatch.customContextLength = committedContext;" in page + + +def test_same_click_commit_covers_all_numeric_inputs(): + """The same-click blur bridge must flush every NumericValueInput-backed + setting, not just Context Length. Max Seq Length (non-GGUF), GPU Layers and + MoE Layers (GGUF) also stage their draft only on blur, so handleRun must + imperatively commit each and fold the value into the staged load config; + otherwise a value the user typed right before clicking Load/Reload is lost.""" + page = _read("features/model-picker/components/model-config-page.tsx") + # Each numeric input owns an imperative handle that handleRun commits, and the + # handle is forwarded down to the actual NumericValueInput. + for ref in ("maxSeqLengthInputRef", "gpuLayersInputRef", "moeLayersInputRef"): + assert f"const {ref} = useRef(null);" in page + assert f"{ref}.current?.commit()" in page + assert f"inputRef={{{ref}}}" in page + # The leaf sub-components accept and forward the handle as a ref. + assert page.count("inputRef?: Ref;") >= 2 + assert "ref={inputRef}" in page + # Committed drafts are folded into the staged config, gated on non-null so an + # untouched field never fabricates an override. + assert "committedMaxSeqLength != null" in page + assert "committedGpuLayers != null" in page + assert "committedMoeLayers != null" in page + assert "pendingPatch.gpuLayers = committedGpuLayers;" in page + assert "pendingPatch.nCpuMoe = committedMoeLayers;" in page + # The non-GGUF load path substitutes the committed Max Seq Length draft. + assert "const effectiveMaxSeqLengthValue =" in page + assert "maxSeqLength: effectiveMaxSeqLengthValue" in page + + +def test_context_commit_rechecks_persistence_only_shortcut(): + """Committed context changes must bypass persistence-only saves.""" + src = _read("features/model-picker/components/model-config-page.tsx") + assert "const effectiveConfig =" in src + assert "perModelConfigsEqual(effectiveConfig, baseline)" in src + assert "const effectivePersistenceOnly =" in src + assert "if (effectivePersistenceOnly)" in src + + def test_reset_enabled_for_explicit_context_pin_at_native(): """An explicit customContextLength that equals the native ceiling is still a user override, so contextAtDefault must require customContextLength == null. From b9d92c41b31e30c65f67fddc33f70eb7ea837d71 Mon Sep 17 00:00:00 2001 From: oobabooga Date: Sat, 25 Jul 2026 07:35:03 -0300 Subject: [PATCH 03/25] Studio: prevent long reasoning from jumping the chat on completion (#7388) --- .../src/components/assistant-ui/reasoning.tsx | 19 ++++++++++++++++++- .../test_chat_response_details_ui_contract.py | 10 ++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/studio/frontend/src/components/assistant-ui/reasoning.tsx b/studio/frontend/src/components/assistant-ui/reasoning.tsx index 97891e9358..09ca6d2530 100644 --- a/studio/frontend/src/components/assistant-ui/reasoning.tsx +++ b/studio/frontend/src/components/assistant-ui/reasoning.tsx @@ -346,6 +346,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ const [manualOpen, setManualOpen] = useState(false); const [dismissedWhileStreaming, setDismissedWhileStreaming] = useState(false); + const [retainStreamingHeight, setRetainStreamingHeight] = useState(false); const [duration, setDuration] = useState(0); const startTimeRef = useRef(null); @@ -368,6 +369,17 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ } }, [isReasoningStreaming]); + // Keep the streaming height cap until the automatic close finishes. Removing + // it on the completion frame expands long reasoning to its full height before + // the collapsible can close, which makes the entire chat jump. + useEffect(() => { + const timeout = window.setTimeout( + () => setRetainStreamingHeight(isReasoningStreaming), + isReasoningStreaming ? 0 : ANIMATION_DURATION, + ); + return () => window.clearTimeout(timeout); + }, [isReasoningStreaming]); + // Open while streaming (unless dismissed), or once manually opened. const isOpen = (isReasoningStreaming && !dismissedWhileStreaming) || manualOpen; const variant = isOpen ? "outline" : "ghost"; @@ -378,6 +390,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ if (isReasoningStreaming) { setDismissedWhileStreaming(!open); } else { + if (open) { + setRetainStreamingHeight(false); + } setManualOpen(open); } }, @@ -407,7 +422,9 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({ aria-busy={isReasoningStreaming} streaming={isReasoningStreaming} > - + {children} diff --git a/tests/studio/test_chat_response_details_ui_contract.py b/tests/studio/test_chat_response_details_ui_contract.py index 8714951883..9d2f0886ee 100644 --- a/tests/studio/test_chat_response_details_ui_contract.py +++ b/tests/studio/test_chat_response_details_ui_contract.py @@ -75,6 +75,16 @@ def test_response_model_badge_is_user_configurable_and_rendered_once_per_message assert 'className="min-w-0 flex-1"' in reasoning_src +def test_reasoning_keeps_streaming_height_cap_through_automatic_collapse(): + src = REASONING_TSX.read_text() + + assert "const [retainStreamingHeight, setRetainStreamingHeight]" in src + assert "setRetainStreamingHeight(false)" in src + assert "setRetainStreamingHeight(isReasoningStreaming)" in src + assert "isReasoningStreaming ? 0 : ANIMATION_DURATION" in src + assert "streaming={isReasoningStreaming || retainStreamingHeight}" in src + + def test_response_details_metadata_is_persisted_without_backend_schema_change(): src = ADAPTER_TS.read_text() assert "interface ResponseDetailsMetadata" in src From 2d026a118430545830370fdf48fc06c5a6acba4d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 25 Jul 2026 04:10:44 -0700 Subject: [PATCH 04/25] Studio: reset quantized KV cache to f16 when the flash-attn-off crash-recovery fallback fires (#7390) * Studio: reset quantized KV cache to f16 when flash-attn-off fallback fires Studio force-enables --flash-attn on for GGUF launches. On a hard startup or first-decode crash it retries via _with_flash_attn_off, which flipped FA off but left --cache-type-k/-v untouched. A quantized KV cache (q8_0, q4_0, q4_1, q5_0, q5_1, iq4_nl) requires flash attention in llama.cpp, so the retry itself aborted at init with 'V cache quantization requires flash_attn' instead of recovering. Reset any quantized --cache-type-k/-v to f16 in the FA-off fallback path so the retry can actually launch. Non-quantized types (f16, bf16, f32) run fine without flash attention and are left unchanged. Handles long and short flag forms and both space and equals syntax, rewriting in place to preserve list length. Adds pytest coverage. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: FA-off fallback resets only the quantized V cache and drops env-only V cache Only the V cache requires flash attention in llama.cpp; a quantized K cache runs fine without it. Restrict the FA-off crash-recovery reset to the V axis (main and draft) so a memory-constrained config keeps its quantized K cache instead of risking an OOM on the recovery. Also drop an inherited quantized V cache set purely through the environment (LLAMA_ARG_CACHE_TYPE_V / LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V) at the FA-off retry sites, which the argv rewrite cannot reach, so the child falls back to the f16 default rather than aborting. * Studio: normalize underscore V-cache aliases in the FA-off fallback llama.cpp rewrites '_' to '-' for any '--' long option before matching, so a pass-through --cache_type_v q8_0 enables a quantized V cache just like --cache-type-v. The FA-off crash-recovery reset only matched the hyphenated spelling, so the underscore alias slipped through and the retry still aborted with "V cache quantization requires flash_attn". Canonicalize the flag name the same way before matching (short flags and the type value are untouched). * [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> --- studio/backend/core/inference/llama_cpp.py | 106 ++++++++++ .../tests/test_llama_cpp_mmproj_fallback.py | 195 ++++++++++++++++++ 2 files changed, 301 insertions(+) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 147174451e..d46e5a0fe1 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -33,6 +33,7 @@ from typing import ( List, Literal, Mapping, + MutableMapping, Optional, Union, ) @@ -3696,6 +3697,14 @@ class LlamaCppBackend: # aborts a --split-mode tensor load, so it's dropped for the tensor attempt. _TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # V cache types that llama.cpp can run WITHOUT flash attention. Only the V + # axis has the dependency: a quantized V cache (q8_0/q4_0/q4_1/q5_0/q5_1/ + # iq4_nl) aborts init with "V cache quantization requires flash_attn", while + # a quantized K cache runs fine without FA. So the flash-attn-off crash- + # recovery fallback must reset a quantized V cache to f16 before it can + # launch (and leaves K alone). These three are the only non-quantized types. + _NON_QUANTIZED_KV_TYPES = frozenset({"f16", "bf16", "f32"}) + # Main-model placement settings that Manual mode owns. They must not leak # from Studio's parent environment into llama-server and silently override # the command assembled from the current request. Draft-model placement is @@ -6149,6 +6158,21 @@ class LlamaCppBackend: cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode) ) + @staticmethod + def _canonical_long_flag(name: str) -> str: + """Return ``name`` with llama.cpp's long-option underscore normalization. + + llama.cpp runs ``std::replace(arg.begin(), arg.end(), '_', '-')`` on any + argv token that starts with ``--`` before looking it up, so a legal + pass-through spelling like ``--cache_type_v`` parses as + ``--cache-type-v``. Mirror that here so managed-flag matching sees the + same canonical name. Short flags (``-ctv``) never carry underscores and + keep their exact spelling; pass only the flag name (no attached value). + """ + if name.startswith("--"): + return name.replace("_", "-") + return name + @staticmethod def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]: """Return cmd with flash attention forced off, or None when its effective @@ -6181,8 +6205,76 @@ class LlamaCppBackend: out[i + 1] = "off" elif explicit(i) is None: # bare flag (reads as on) -> explicit off out[i] = f"{tok}=off" + + # A quantized V cache requires flash attention in llama.cpp: the init + # aborts with "V cache quantization requires flash_attn". A quantized K + # cache has no such requirement and runs fine without FA, so it is left + # untouched -- resetting it would needlessly enlarge the K cache and can + # OOM a memory-constrained config. Studio launches with FA on, so a + # quantized --cache-type-v is legal at launch but would make THIS FA-off + # retry crash on init instead of recovering. Reset a quantized V cache -- + # main and draft (the draft context shares the global --flash-attn flag, + # so its V cache aborts too) -- to f16 (the llama.cpp default); + # non-quantized types -- f16/bf16/f32 -- run fine without FA and are left + # untouched. The value is rewritten in place so the list length is + # preserved for downstream slices, matching the flash-attn flip above. + _v_cache_flags = ( + "--cache-type-v", + "-ctv", + "--cache-type-v-draft", + "--spec-draft-type-v", + "-ctvd", + ) + _cache_reset = False + for i, tok in enumerate(out): + # llama.cpp rewrites '_' to '-' for any argv token starting with + # '--' before matching, so a legal pass-through spelling such as + # --cache_type_v parses as --cache-type-v and still enables a + # quantized V cache. Canonicalize the flag name the same way so the + # reset recognizes the underscore aliases too; short flags (-ctv) + # and the type value are left untouched. + name = LlamaCppBackend._canonical_long_flag(tok.partition("=")[0]) + if name not in _v_cache_flags: + continue + if "=" in tok: + flag, _, value = tok.partition("=") + if value.strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i] = f"{flag}=f16" + _cache_reset = True + elif i + 1 < len(out): + if out[i + 1].strip().lower() not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + out[i + 1] = "f16" + _cache_reset = True + if _cache_reset: + logger.info( + "V cache dtype reset to f16 because flash attention was disabled " + "by the crash-recovery fallback (quantized V cache requires flash " + "attention in llama.cpp; the K cache is left untouched)." + ) return out + @staticmethod + def _drop_env_quantized_v_cache(env: MutableMapping[str, str]) -> bool: + """Drop an inherited quantized V-cache env var (main or draft) in place + before a flash-attn-off retry, returning True if anything was removed. + + The argv rewrite in ``_with_flash_attn_off`` only reaches flags on the + command line. Studio deliberately lets an env-only cache type reach the + child untouched (an asymmetric K/V env must survive), so a quantized V + cache set purely through ``LLAMA_ARG_CACHE_TYPE_V`` (or the draft + ``LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V``) would still abort the FA-off retry + with "V cache quantization requires flash_attn". Dropping it lets + llama.cpp fall back to the f16 default. Only V is dropped: a quantized K + cache runs fine without flash attention, so its env var is preserved. + """ + dropped = False + for var in ("LLAMA_ARG_CACHE_TYPE_V", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V"): + value = (env.get(var) or "").strip().lower() + if value and value not in LlamaCppBackend._NON_QUANTIZED_KV_TYPES: + env.pop(var, None) + dropped = True + return dropped + @staticmethod def _strip_mmproj_args(cmd: list[str]) -> list[str]: """Return cmd without the '--mmproj ' pair (text-only retry). @@ -8339,6 +8431,13 @@ class LlamaCppBackend: _fa_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = _spawn_and_wait(_fa_cmd, label = "-noflash") @@ -8384,6 +8483,13 @@ class LlamaCppBackend: _probe_rc, ) self._kill_process() + # The argv rewrite can't reach an env-only quantized V + # cache; drop it so the FA-off child doesn't abort on it. + if self._drop_env_quantized_v_cache(env): + logger.info( + "Dropped inherited quantized V-cache env for the " + "--flash-attn off retry (requires flash attention)." + ) cmd = _fa_cmd healthy = ( _spawn_and_wait(_fa_cmd, label = "-noflash-mtp") diff --git a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py index 4332a440a5..45c8bcb032 100644 --- a/studio/backend/tests/test_llama_cpp_mmproj_fallback.py +++ b/studio/backend/tests/test_llama_cpp_mmproj_fallback.py @@ -250,6 +250,201 @@ class TestFlashAttnOff: assert _flash_off(["llama-server", "-fa"]) == ["llama-server", "-fa=off"] +_drop_env_v = LlamaCppBackend._drop_env_quantized_v_cache + + +class TestFlashAttnOffQuantizedKvCache: + """Only the V cache requires flash attention in llama.cpp (init aborts with + "V cache quantization requires flash_attn"); a quantized K cache runs fine + without FA. Studio launches FA on, so a quantized --cache-type-v is legal at + launch but would make the FA-off crash-recovery retry crash on init. The + fallback must reset a quantized V cache (main and draft) to f16 while leaving + the K cache and non-quantized (f16/bf16/f32) types unchanged -- resetting K + would needlessly enlarge it and can OOM a memory-constrained config.""" + + _QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"] + _NON_QUANTIZED = ["f16", "bf16", "f32"] + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_quantized_v_reset_k_preserved(self, qtype): + cmd = [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + qtype, + "--cache-type-v", + qtype, + ] + out = _flash_off(cmd) + assert out is not None + # FA flipped off AND the V axis reset to f16; the K axis is preserved so + # the FA-off retry keeps its memory budget (quantized K is FA-independent). + assert out[out.index("--flash-attn") + 1] == "off" + assert out[out.index("--cache-type-k") + 1] == qtype + assert out[out.index("--cache-type-v") + 1] == "f16" + assert len(out) == len(cmd) + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_quantized_draft_v_reset(self, qtype): + # The draft context shares the global --flash-attn flag, so its quantized + # V cache aborts too and must be reset; the draft K cache is preserved. + for v_flag, k_flag in ( + ("--cache-type-v-draft", "--cache-type-k-draft"), + ("--spec-draft-type-v", "--spec-draft-type-k"), + ("-ctvd", "-ctkd"), + ): + cmd = ["llama-server", "-fa", "on", k_flag, qtype, v_flag, qtype] + out = _flash_off(cmd) + assert out is not None + assert out[out.index(v_flag) + 1] == "f16" + assert out[out.index(k_flag) + 1] == qtype + + @pytest.mark.parametrize("ntype", _NON_QUANTIZED) + def test_nonquantized_cache_left_unchanged(self, ntype): + cmd = [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + ntype, + "--cache-type-v", + ntype, + ] + out = _flash_off(cmd) + assert out is not None + # Only FA flips; the non-quantized cache type is preserved verbatim. + assert out[out.index("--flash-attn") + 1] == "off" + assert out[out.index("--cache-type-k") + 1] == ntype + assert out[out.index("--cache-type-v") + 1] == ntype + + def test_equals_form_quantized_v_reset(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-v=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache-type-v=f16"] + + def test_equals_form_quantized_k_preserved(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache-type-k=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache-type-k=q8_0"] + + def test_short_alias_v_reset_k_preserved(self): + out = _flash_off(["llama-server", "-fa", "on", "-ctk", "q4_0", "-ctv", "q4_0"]) + assert out == ["llama-server", "-fa", "off", "-ctk", "q4_0", "-ctv", "f16"] + + def test_asymmetric_cache_only_v_reset(self): + # Quantized V, non-quantized K: reset V, keep K untouched. + out = _flash_off( + [ + "llama-server", + "--flash-attn", + "on", + "--cache-type-k", + "f16", + "--cache-type-v", + "q8_0", + ] + ) + assert out[out.index("--cache-type-k") + 1] == "f16" + assert out[out.index("--cache-type-v") + 1] == "f16" + + def test_no_cache_flags_still_flips_fa(self): + out = _flash_off(["llama-server", "--flash-attn", "on", "-c", "4096"]) + assert out == ["llama-server", "--flash-attn", "off", "-c", "4096"] + + def test_quantized_k_only_still_flips_fa_but_keeps_k(self): + # A quantized K cache with no V flag is a valid FA-off launch; the retry + # must not touch the K cache (it would waste memory for nothing). + out = _flash_off(["llama-server", "--flash-attn", "on", "--cache-type-k", "q8_0"]) + assert out == ["llama-server", "--flash-attn", "off", "--cache-type-k", "q8_0"] + + def test_input_not_mutated(self): + cmd = ["llama-server", "--flash-attn", "on", "--cache-type-v", "q8_0"] + _flash_off(cmd) + assert cmd[-1] == "q8_0" + + @pytest.mark.parametrize( + "flag", + ["--cache_type_v", "--cache-type_v", "--cache_type-v"], + ) + def test_underscore_alias_v_reset(self, flag): + # llama.cpp normalizes '_' to '-' in any '--' long option before + # matching, so a pass-through --cache_type_v enables a quantized V cache + # and must be reset by the FA-off retry too (else init aborts). + out = _flash_off(["llama-server", "--flash-attn", "on", flag, "q8_0"]) + assert out is not None + assert out[out.index("--flash-attn") + 1] == "off" + # The user's flag spelling is preserved; llama.cpp normalizes it anyway. + assert out[out.index(flag) + 1] == "f16" + + def test_underscore_alias_draft_v_reset(self): + out = _flash_off(["llama-server", "-fa", "on", "--spec_draft_type_v", "q4_0"]) + assert out is not None + assert out[out.index("--spec_draft_type_v") + 1] == "f16" + + def test_underscore_alias_equals_form_v_reset(self): + out = _flash_off(["llama-server", "--flash-attn=on", "--cache_type_v=q8_0"]) + assert out == ["llama-server", "--flash-attn=off", "--cache_type_v=f16"] + + def test_underscore_value_not_normalized_for_nonquantized(self): + # Only the flag name is canonicalized; a non-quantized type value is + # matched verbatim and left untouched (no spurious reset). + out = _flash_off(["llama-server", "--flash-attn", "on", "--cache_type_v", "f16"]) + assert out[out.index("--cache_type_v") + 1] == "f16" + assert out[out.index("--flash-attn") + 1] == "off" + + def test_short_alias_underscore_not_applied(self): + # Short flags are never underscore-normalized by llama.cpp; -ctv still + # matches and resets, and an unrelated short token is left alone. + out = _flash_off(["llama-server", "-fa", "on", "-ctv", "q8_0"]) + assert out == ["llama-server", "-fa", "off", "-ctv", "f16"] + + +class TestDropEnvQuantizedVCache: + """The argv rewrite can't reach a cache type set purely through the + environment (Studio deliberately lets an env-only type reach the child), so + the FA-off retry separately drops a quantized V-cache env var. Only V is + dropped: a quantized K cache is FA-independent and must survive.""" + + _QUANTIZED = ["q8_0", "q4_0", "q4_1", "q5_0", "q5_1", "iq4_nl"] + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_drops_quantized_main_v_env(self, qtype): + env = {"LLAMA_ARG_CACHE_TYPE_V": qtype, "PATH": "/usr/bin"} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_CACHE_TYPE_V" not in env + assert env["PATH"] == "/usr/bin" + + @pytest.mark.parametrize("qtype", _QUANTIZED) + def test_drops_quantized_draft_v_env(self, qtype): + env = {"LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V": qtype} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_V" not in env + + def test_preserves_quantized_k_env(self): + # A quantized K cache runs without FA, so its env must not be dropped. + env = {"LLAMA_ARG_CACHE_TYPE_K": "q8_0", "LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K": "q4_0"} + assert _drop_env_v(env) is False + assert env["LLAMA_ARG_CACHE_TYPE_K"] == "q8_0" + assert env["LLAMA_ARG_SPEC_DRAFT_CACHE_TYPE_K"] == "q4_0" + + @pytest.mark.parametrize("ntype", ["f16", "bf16", "f32", "F16", " q8_0 "]) + def test_preserves_nonquantized_v_env(self, ntype): + # Non-quantized V env values (and whitespace/case variants of them) run + # fine without FA; only a genuinely quantized value is dropped. + if ntype.strip().lower() in ("q8_0",): + env = {"LLAMA_ARG_CACHE_TYPE_V": ntype} + assert _drop_env_v(env) is True + assert "LLAMA_ARG_CACHE_TYPE_V" not in env + else: + env = {"LLAMA_ARG_CACHE_TYPE_V": ntype} + assert _drop_env_v(env) is False + assert env["LLAMA_ARG_CACHE_TYPE_V"] == ntype + + def test_noop_on_empty_env(self): + env = {} + assert _drop_env_v(env) is False + assert env == {} + + class TestNonProjectorDiagnostic: """_output_has_nonprojector_diagnostic gates the signal-only text-only retry: a hard crash that already names OOM / a bad arch / a TP limit must surface From c4b777263dbb855b2f3ef8cd9cd480a4b191973a Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 25 Jul 2026 04:11:03 -0700 Subject: [PATCH 05/25] fix(studio/colab): fix OutStream startup crash and tidy the notebook cards (#7404) * fix(studio/colab): survive ipykernel OutStream close() during startup Unsloth Studio crashed at server startup on Colab with: Unsloth Studio failed to start: 'OutStream' object has no attribute 'watch_fd_thread' Root cause: - Colab's ipykernel OutStream is created with watchfd=False, so it never gains a watch_fd_thread. The OutStream.close() in the affected ipykernel versions joins that thread unconditionally and raises AttributeError (ipython/ipykernel#867). - _setup_server_disk_logging() replaces sys.stdout/sys.stderr with a tee. That changes the console object identity, so Colab's absl logging handler (which captured the original OutStream and whose close() deliberately skips sys.stdout/sys.stderr) no longer treats it as the live console. - run_server builds uvicorn.Config(...), whose configure_logging runs logging.config.dictConfig -> logging.shutdown, closing every existing handler. The absl handler then calls close() on the orphaned OutStream and the AttributeError propagates out of uvicorn.Config and aborts startup. Fix: - Before installing the tee, harden the displaced console streams' close() so only the ipykernel#867 AttributeError is swallowed; a healthy close() runs unchanged and any other error still propagates. The buggy close() raises before it nulls pub_thread, so the stream stays fully usable. - Give _TeeStream its own close() that flushes the log copy and forwards close() to the wrapped console stream best-effort, so a handler that captured the tee cannot crash startup either. Add regression tests reproducing the exact path (an absl-style handler closing a watchfd=False OutStream stand-in during logging.shutdown) and asserting the tee/console path survives and keeps logging. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Show the Colab login password in the shareable link card * Tighten Colab card comments for PR #7404 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make the Colab tunnel URL clickable and emphasise the password * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Narrow the console close() hardening to the watch_fd_thread AttributeError * Put the Colab password on its own line so selection excludes the label * Keep the Colab password as plain selectable text --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- studio/backend/colab.py | 68 +++-- studio/backend/run.py | 83 ++++++ studio/backend/tests/test_colab_embed.py | 125 ++++++++- .../test_server_disk_logging_outstream.py | 258 ++++++++++++++++++ 4 files changed, 512 insertions(+), 22 deletions(-) create mode 100644 studio/backend/tests/test_server_disk_logging_outstream.py diff --git a/studio/backend/colab.py b/studio/backend/colab.py index baa18a2fec..051d80abfe 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -218,11 +218,10 @@ def _colab_login_html(username: str, password: str) -> str: Unsloth Studio Login (Colab)

- Log in to Studio with the Cloudflare link above using these credentials. This cell - is visible only in your notebook session. + Log in as {username} with this password. This cell is visible only in + your notebook session.

- Username: {username}
Password: {password}

@@ -441,8 +440,29 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: return False -def _shareable_link_html(cloudflare_url: str) -> str: - """Branded card for the shareable Cloudflare link, styled like the show_link banner.""" +def _shareable_link_html( + cloudflare_url: str, + password: "str | None" = None, + username: "str | None" = None, +) -> str: + """Branded card for the shareable Cloudflare link, styled like the show_link banner. + + *password* renders under the link so the credential sits in the card with the button + it unlocks. The username is always the default admin, so it reads inline. + """ + login_block = "" + if password: + login_block = f""" +

+ Password +

+

{password}

+

+ Log in as {username} with this password. Shown only in your + notebook session, and never included in the shared link. +

""" return f"""
@@ -460,11 +480,12 @@ def _shareable_link_html(cloudflare_url: str) -> str: Open Unsloth Studio

- This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab. + This Cloudflare HTTPS link works from any device, so you can share it with anyone.

- 🔗 {cloudflare_url} -

+ 🔗 {cloudflare_url} +

{login_block}
""" @@ -555,28 +576,37 @@ def _show_and_embed( cloudflare_url = cloudflare_url, ) + # Fold the credentials into the link card rather than a second card below it. + credentials_shown = False if cloudflare_url: try: from IPython.display import HTML, display - display(HTML(_shareable_link_html(cloudflare_url))) + + username, password = colab_login if colab_login else (None, None) + display(HTML(_shareable_link_html(cloudflare_url, password, username))) + credentials_shown = bool(colab_login) except Exception as e: logger.info(f"Could not render Cloudflare link card ({e}).") - if colab_login: + if colab_login and not credentials_shown: try: _show_colab_login_credentials(*colab_login) except Exception as e: logger.info(f"Could not render Colab login card ({e}).") - try: - show_link( - port, - _url = url, - has_cloudflare_link = bool(cloudflare_url), - cloudflare_requested = cloudflare_requested, - ) - except Exception as e: - logger.info(f"Could not render Unsloth link card ({e}).") + # With a tunnel up the embed below is skipped, so the ready card would only restate + # the link card and print a proxy URL that 404s outside this tab. + skip_ready_card = _is_colab_runtime() and bool(cloudflare_url) + if not skip_ready_card: + try: + show_link( + port, + _url = url, + has_cloudflare_link = bool(cloudflare_url), + cloudflare_requested = cloudflare_requested, + ) + except Exception as e: + logger.info(f"Could not render Unsloth link card ({e}).") # On Colab with a working tunnel, skip the in-cell proxy embed (often blank). if _is_colab_runtime() and cloudflare_url: diff --git a/studio/backend/run.py b/studio/backend/run.py index d9569c46f6..90fd28670c 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -991,10 +991,88 @@ class _TeeStream: except Exception: pass + def close(self): + # We do NOT own the console stream (it is the terminal / Jupyter kernel + # stream we wrapped), so closing the tee must never take the server down. + # Flush the log copy, then forward close() to the wrapped stream + # best-effort: on Colab that stream is an ipykernel OutStream whose + # close() can raise (see _harden_console_close / ipython/ipykernel#867). + try: + self._log_fh.flush() + except Exception: + pass + try: + self._stream.close() + except Exception: + pass + def __getattr__(self, name): return getattr(self._stream, name) +_WATCH_FD_THREAD_ATTR = "watch_fd_thread" + + +def _is_missing_watch_fd_thread(exc): + """True only for ipython/ipykernel#867's missing-``watch_fd_thread`` error. + + ``AttributeError.name`` exists from Python 3.10; the message carries the + attribute name on every version (possibly with a "Did you mean" tail), so + check both and let every other AttributeError through. + """ + if getattr(exc, "name", None) == _WATCH_FD_THREAD_ATTR: + return True + return _WATCH_FD_THREAD_ATTR in str(exc) + + +def _harden_console_close(stream): + """Stop a displaced console stream's close() from aborting Studio startup. + + ``_setup_server_disk_logging`` replaces ``sys.stdout``/``sys.stderr`` with a + tee. That changes the object identity of the console stream, so a third-party + logging handler that captured the ORIGINAL stream (notably Colab's ``absl`` + logging handler, whose ``close()`` skips ``sys.stdout``/``sys.stderr`` but not + a stream that is no longer either) treats it as an ordinary stream and calls + ``close()`` on it during logging teardown -- ``uvicorn.Config()`` -> + ``logging.config.dictConfig()`` -> ``logging.shutdown()``. + + A Jupyter/Colab ``ipykernel`` ``OutStream`` created with ``watchfd=False`` + (the Colab default, and every in-process kernel) never gains a + ``watch_fd_thread``, yet the ``OutStream.close()`` shipped in the affected + ipykernel versions joins that thread unconditionally and raises + ``AttributeError: 'OutStream' object has no attribute 'watch_fd_thread'`` + (ipython/ipykernel#867). That AttributeError propagates out of + ``uvicorn.Config(...)`` and aborts startup ("Unsloth Studio failed to start"). + + Wrap the stream's ``close()`` in a transparent pass-through that swallows + ONLY that specific teardown AttributeError. A healthy close() (a real console + stream, or an OutStream with fd-watching on) runs to completion exactly as + before and any other error still propagates, so nothing changes off Colab. A + stream whose ``close`` cannot be reassigned keeps its original close(). + """ + try: + _orig_close = stream.close + except Exception: + return + + def _safe_close(*args, **kwargs): + try: + return _orig_close(*args, **kwargs) + except AttributeError as exc: + if not _is_missing_watch_fd_thread(exc): + # A real teardown failure; never hide it. + raise + # ipython/ipykernel#867: watchfd=False OutStream.close() joins a + # thread that was never created. Nothing to clean up; keep going. + return None + + try: + stream.close = _safe_close + except (AttributeError, TypeError): + # A stream that forbids setting instance attributes; leave it as-is. + pass + + def _setup_server_disk_logging(): """Tee stdout/stderr to ~/.unsloth/studio/logs/server/ and aim faulthandler at the same file so hard crashes (access violations / @@ -1037,6 +1115,11 @@ def _setup_server_disk_logging(): # the stderr the server already captures. os.environ.setdefault("PYTHONFAULTHANDLER", "1") + # Replacing the console streams orphans them from third-party "is this the + # live console?" checks, so guard their close() first (ipython/ipykernel#867). + _harden_console_close(sys.stdout) + _harden_console_close(sys.stderr) + sys.stdout = _TeeStream(sys.stdout, log_fh) sys.stderr = _TeeStream(sys.stderr, log_fh) diff --git a/studio/backend/tests/test_colab_embed.py b/studio/backend/tests/test_colab_embed.py index dae0c7dae0..83b2a5a82d 100644 --- a/studio/backend/tests/test_colab_embed.py +++ b/studio/backend/tests/test_colab_embed.py @@ -336,17 +336,71 @@ def test_colab_login_html_includes_credentials(): html = colab._colab_login_html("unsloth", "alpha-beta-gamma-delta") assert "unsloth" in html assert "alpha-beta-gamma-delta" in html + # The username is fixed, so it reads inline rather than as its own field. + assert "Username:" not in html -def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch): +def test_shareable_link_html_embeds_password_under_the_link(): + """The credential belongs in the same card as the button it unlocks.""" + html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth") + assert "share.trycloudflare.com" in html + assert "secret-pass" in html + # Username is stated inline, not as its own labelled field. + assert "Username:" not in html + assert "unsloth" in html + # The password must sit after the link, not above it. + assert html.index("share.trycloudflare.com") < html.index("secret-pass") + + +def test_shareable_link_html_renders_the_url_as_a_link(): + """The printed URL is an anchor, using the popup-safe open the button uses.""" + html = colab._shareable_link_html("https://share.trycloudflare.com") + assert 'https://share.trycloudflare.com" in html + assert html.count("window.open(this.href,'_blank')") == 2 + + +def test_shareable_link_html_emphasises_the_password(): + """The password is the one thing to copy, so it is enlarged and underlined.""" + html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth") + pw_tag = html[html.index("Password") : html.index("secret-pass")] + assert "font-size: 24px" in pw_tag + assert "text-decoration: underline" in pw_tag + + +def test_shareable_link_html_password_has_no_adjacent_whitespace(): + """Whitespace beside the password is selected with it on a double click.""" + html = colab._shareable_link_html("https://share.trycloudflare.com", "secret-pass", "unsloth") + before, after = html.split("secret-pass", 1) + assert before.endswith(">") + assert after.startswith("<") + # Label on its own line, so nothing shares the password's text node. + assert "Password:" not in html + # Plain selectable text: user-select overrides break double click to select. + assert "user-select" not in html + + +def test_shareable_link_html_omits_login_block_without_password(): + html = colab._shareable_link_html("https://share.trycloudflare.com") + assert "Password" not in html + + +def test_show_and_embed_folds_login_into_the_cloudflare_card(monkeypatch): + """One card, not two: the tunnel card carries the password itself.""" displayed: list[str] = [] ipython_display = SimpleNamespace( HTML = lambda html: SimpleNamespace(html = html), display = lambda html: displayed.append(html.html), ) + login_cards: list[tuple] = [] monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_show_colab_login_credentials", + lambda *args: login_cards.append(args), + ) monkeypatch.setattr( colab, "show_link", @@ -360,9 +414,74 @@ def test_show_and_embed_renders_cloudflare_before_colab_login(monkeypatch): colab_login = ("unsloth", "secret-pass"), ) - assert len(displayed) == 2 + assert len(displayed) == 1 assert "share.trycloudflare.com" in displayed[0] - assert "secret-pass" in displayed[1] + assert "secret-pass" in displayed[0] + assert login_cards == [] + + +def test_show_and_embed_keeps_separate_login_card_without_tunnel(monkeypatch): + """No tunnel card to fold into, so the standalone login card still renders.""" + login_cards: list[tuple] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "_show_colab_login_credentials", + lambda *args: login_cards.append(args), + ) + monkeypatch.setattr( + colab, + "show_link", + lambda port, *, _url = None, has_cloudflare_link = False, cloudflare_requested = False: None, + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + colab._show_and_embed(8888, colab_login = ("unsloth", "secret-pass")) + + assert login_cards == [("unsloth", "secret-pass")] + + +def test_show_and_embed_skips_ready_card_when_tunnel_is_up(monkeypatch): + """The ready card only restates the tunnel card and prints a proxy URL that 404s.""" + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + colab._show_and_embed(8888, cloudflare_url = "https://share.trycloudflare.com") + + assert calls == [] + + +def test_show_and_embed_keeps_ready_card_without_tunnel(monkeypatch): + """Without a tunnel the ready card is the only guidance, so it must stay.""" + calls: list[str] = [] + + monkeypatch.setattr(colab, "get_colab_url", lambda port: "https://8888-test.prod.colab.dev/") + monkeypatch.setattr(colab, "_is_colab_runtime", lambda: True) + monkeypatch.setattr( + colab, + "show_link", + lambda port, + *, + _url = None, + has_cloudflare_link = False, + cloudflare_requested = False: calls.append("show_link"), + ) + monkeypatch.setattr(colab, "_embed_kernel_port_iframe", lambda port: True) + colab._show_and_embed(8888) + + assert calls == ["show_link"] def test_show_and_embed_skips_iframe_on_colab_when_cloudflare_ready(monkeypatch): diff --git a/studio/backend/tests/test_server_disk_logging_outstream.py b/studio/backend/tests/test_server_disk_logging_outstream.py new file mode 100644 index 0000000000..0ff27666a0 --- /dev/null +++ b/studio/backend/tests/test_server_disk_logging_outstream.py @@ -0,0 +1,258 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Regression tests for the Colab "OutStream has no attribute 'watch_fd_thread'" +startup crash. + +Field report (Colab): Unsloth Studio dies at server startup with +``❌ Unsloth Studio failed to start: 'OutStream' object has no attribute +'watch_fd_thread'``. + +Root cause chain: + * Colab's ipykernel ``OutStream`` is created with ``watchfd=False``, so it + never gains a ``watch_fd_thread``; the ``OutStream.close()`` shipped in the + affected ipykernel versions joins that thread unconditionally and raises + ``AttributeError`` (ipython/ipykernel#867). + * ``run._setup_server_disk_logging()`` replaces ``sys.stdout``/``sys.stderr`` + with a ``_TeeStream``. That changes the console object identity, so Colab's + ``absl`` logging handler -- which captured the ORIGINAL OutStream and whose + ``close()`` deliberately skips ``sys.stdout``/``sys.stderr`` -- no longer + recognizes it as the live console. + * ``run_server`` builds ``uvicorn.Config(...)``, whose ``configure_logging`` -> + ``logging.config.dictConfig`` -> ``logging.shutdown`` closes every existing + handler. The absl handler then calls ``OutStream.close()`` on the orphaned + stream, and the AttributeError aborts startup. + +These tests reproduce the mechanism with a stand-in OutStream (Colab-identical +constructs are not importable off Colab) and assert the tee/console path used at +startup survives it. +""" + +from __future__ import annotations + +import io +import logging +import sys +import weakref +from pathlib import Path + +import pytest + +_BACKEND_DIR = str(Path(__file__).resolve().parent.parent) +if _BACKEND_DIR not in sys.path: + sys.path.insert(0, _BACKEND_DIR) + +import run as run_mod # noqa: E402 + + +class _ColabOutStream(io.TextIOBase): + """Stand-in for Colab's ipykernel OutStream built with ``watchfd=False``: + no ``watch_fd_thread`` and an unguarded ``close()`` that joins it + (ipython/ipykernel#867).""" + + def __init__(self, name: str, sink: io.StringIO): + self.name = name + self._sink = sink + + def write(self, s): + return self._sink.write(s) + + def flush(self): + pass + + def writable(self): + return True + + def isatty(self): + return False + + def close(self): + # Never set because watchfd=False -> AttributeError, exactly as Colab. + self.watch_fd_thread.join() + + def __del__(self): + # io.TextIOBase.__del__ would call our buggy close() at GC (the harmless + # "Exception ignored" tail seen in Colab); silence it so the test is clean. + pass + + +class _WatchingOutStream(_ColabOutStream): + """OutStream with fd-watching ON: ``watch_fd_thread`` exists, close() is + well behaved and must keep working unchanged.""" + + def __init__(self, name: str, sink: io.StringIO): + super().__init__(name, sink) + self.close_ran = False + self.watch_fd_thread = type("_T", (), {"join": lambda self: None})() + + def close(self): + self.watch_fd_thread.join() + self.close_ran = True + + +class _AbslLikeHandler(logging.StreamHandler): + """Mirror of ``absl.logging.PythonHandler.close()``: close the captured + stream unless it is (still) one of the user-managed console streams.""" + + def close(self): + try: + user_managed = (sys.stderr, sys.stdout, sys.__stderr__, sys.__stdout__) + if self.stream not in user_managed and ( + not hasattr(self.stream, "isatty") or not self.stream.isatty() + ): + self.stream.close() + except ValueError: + pass + super().close() + + +class TestHardenConsoleClose: + def test_neutralizes_watchfd_false_close(self): + stream = _ColabOutStream("stdout", io.StringIO()) + with pytest.raises(AttributeError): + stream.close() # baseline: the ipykernel #867 bug is real + + stream = _ColabOutStream("stdout", io.StringIO()) + run_mod._harden_console_close(stream) + assert stream.close() is None # swallowed, no crash + + def test_healthy_close_still_runs_fully(self): + stream = _WatchingOutStream("stdout", io.StringIO()) + run_mod._harden_console_close(stream) + stream.close() + assert stream.close_ran is True + + def test_only_attributeerror_is_swallowed(self): + class _Boom: + def close(self): + raise ValueError("real teardown failure") + + stream = _Boom() + run_mod._harden_console_close(stream) + with pytest.raises(ValueError): + stream.close() + + def test_unrelated_attributeerror_still_propagates(self): + # Only #867 is neutralized; a genuine missing attribute during teardown + # must still surface instead of looking like a clean close. + class _Console: + def close(self): + return self.not_a_real_attribute + + stream = _Console() + run_mod._harden_console_close(stream) + with pytest.raises(AttributeError, match = "not_a_real_attribute"): + stream.close() + + def test_swallowed_across_attributeerror_message_shapes(self): + # Python 3.12 appends a "Did you mean" tail; the match must survive it, + # and pre-3.10 AttributeErrors carry no ``name``, only the message. + class _Suggesting: + def close(self): + raise AttributeError( + "'OutStream' object has no attribute 'watch_fd_thread'. " + "Did you mean: '_watch_pipe_fd'?" + ) + + stream = _Suggesting() + run_mod._harden_console_close(stream) + assert stream.close() is None + + def test_unsettable_close_is_left_alone(self): + # A stream whose close cannot be reassigned must not raise from hardening. + class _Frozen: + __slots__ = () + + def close(self): + return "ok" + + stream = _Frozen() + run_mod._harden_console_close(stream) # must not raise + assert stream.close() == "ok" + + +class TestTeeStreamClose: + def test_tee_close_over_buggy_stream_never_raises(self): + console = _ColabOutStream("stdout", io.StringIO()) + log = io.StringIO() + tee = run_mod._TeeStream(console, log) + tee.write("before-close") + tee.close() # must not raise despite the wrapped stream's broken close + assert log.getvalue() == "before-close" + + def test_tee_close_flushes_log(self): + class _FlushCounting(io.StringIO): + def __init__(self): + super().__init__() + self.flushes = 0 + + def flush(self): + self.flushes += 1 + super().flush() + + console, log = io.StringIO(), _FlushCounting() + tee = run_mod._TeeStream(console, log) + tee.write("x") + tee.close() + assert log.flushes >= 1 + + +class TestColabStartupRegression: + """End-to-end: the exact trigger -- an absl-style handler closing the + orphaned OutStream during the ``logging.shutdown`` that uvicorn's + ``uvicorn.Config`` -> ``dictConfig`` runs -- must not crash Studio, and the + tee must keep logging afterwards. + + ``logging.shutdown`` is driven over a LOCAL weakref list (identical code path + to ``logging.config._clearExistingHandlers``) so the global logging state and + pytest's own capture are untouched. + """ + + def _make_console_and_handlers(self, monkeypatch): + out_sink, err_sink = io.StringIO(), io.StringIO() + out_stream = _ColabOutStream("stdout", out_sink) + err_stream = _ColabOutStream("stderr", err_sink) + monkeypatch.setattr(sys, "stdout", out_stream) + monkeypatch.setattr(sys, "stderr", err_stream) + # absl-like handlers capture the ORIGINAL OutStreams (as in Colab). + handlers = [_AbslLikeHandler(sys.stdout), _AbslLikeHandler(sys.stderr)] + return out_sink, err_sink, out_stream, err_stream, handlers + + def test_baseline_reproduces_crash_without_fix(self, monkeypatch): + # Prove the test exercises the real path: swapping the console identity + # (what the tee does) makes the absl-like close hit #867. + _, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch) + try: + monkeypatch.setattr(sys, "stdout", io.StringIO()) + monkeypatch.setattr(sys, "stderr", io.StringIO()) + with pytest.raises(AttributeError, match = "watch_fd_thread"): + logging.shutdown([weakref.ref(h) for h in handlers]) + finally: + # Neutralize so a lingering handler can't crash global teardown. + run_mod._harden_console_close(out_stream) + run_mod._harden_console_close(err_stream) + for h in handlers: + try: + h.close() + except Exception: + pass + + def test_startup_survives_with_harden_and_tee(self, monkeypatch): + out_sink, _, out_stream, err_stream, handlers = self._make_console_and_handlers(monkeypatch) + + # Exactly what _setup_server_disk_logging does before serving: + run_mod._harden_console_close(sys.stdout) + run_mod._harden_console_close(sys.stderr) + log_fh = io.StringIO() + monkeypatch.setattr(sys, "stdout", run_mod._TeeStream(sys.stdout, log_fh)) + monkeypatch.setattr(sys, "stderr", run_mod._TeeStream(sys.stderr, log_fh)) + + # The close-storm uvicorn triggers via dictConfig -> logging.shutdown, + # closing the absl-like handlers over the (now orphaned) OutStreams. + logging.shutdown([weakref.ref(h) for h in handlers]) # must NOT raise + + # The tee still tees to both console and disk afterwards. + print("post-startup-line") + sys.stdout.flush() + assert "post-startup-line" in out_sink.getvalue() + assert "post-startup-line" in log_fh.getvalue() From 85f6231a2fb48d745a9ebaf5792d9c6916740e1f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sat, 25 Jul 2026 04:42:33 -0700 Subject: [PATCH 06/25] tests: anchor the gguf ordering assertion on the branch that owns the marker (#7443) _load_model_impl contains more than one `if config.is_gguf:`, so source.index() returned the earlier one, which belongs to a different check than the branch the assertion is reasoning about. The inheritance call sits at line 4543, the earlier branch at 4508 and the branch holding the load marker at 4567, so the comparison read 186995 < 185014 and failed on main. The branch is now located from the load marker itself, which is the landmark the rest of the test already relies on, so the assertion compares the inheritance call against the branch that actually guards it. The slice used by the following assertions is anchored the same way, which also tightens them: they previously searched from the earlier branch to end of file. The invariant is unchanged and still has teeth: moving the inheritance call after the branch makes the assertion fail. Co-authored-by: danielhanchen --- studio/backend/tests/test_gguf_load_cache_reuse.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/studio/backend/tests/test_gguf_load_cache_reuse.py b/studio/backend/tests/test_gguf_load_cache_reuse.py index dfd6fc6034..ce26147b11 100644 --- a/studio/backend/tests/test_gguf_load_cache_reuse.py +++ b/studio/backend/tests/test_gguf_load_cache_reuse.py @@ -785,7 +785,12 @@ class TestLoadHubDownloadExclusion: def test_load_marker_precedes_hub_guard_and_unload(self): source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text() - gguf_branch = source[source.index("if config.is_gguf:") :] + # _load_model_impl has more than one `if config.is_gguf:`, so anchor on + # the branch that actually owns the load marker rather than the first + # one in the file, which belongs to an earlier check. + marker = source.index("enter_context(gguf_load_in_flight") + gguf_branch_start = source.rindex("if config.is_gguf:", 0, marker) + gguf_branch = source[gguf_branch_start:] # The gguf_load_in_flight marker must be entered before the hub-download # guard and the unload so a concurrent load can't race the download @@ -794,7 +799,7 @@ class TestLoadHubDownloadExclusion: # inherited value (e.g. a carried --no-mmproj) shapes the guard's # require_mmproj. Anchor on the call form so the assertion pins the # endpoint's call site, not the function definition. - assert source.index("= _resolve_inherited_extra_args(") < source.index("if config.is_gguf:") + assert source.index("= _resolve_inherited_extra_args(") < gguf_branch_start assert ( gguf_branch.index("enter_context(gguf_load_in_flight") < gguf_branch.index("_hub_download_blocks_gguf_load") From 3ea6d14c395a932ad97ada16ca37c5af675251e8 Mon Sep 17 00:00:00 2001 From: Leo Borcherding Date: Sat, 25 Jul 2026 18:58:02 -0500 Subject: [PATCH 07/25] AMD: CI coverage for recent fixes, plus three wrong gfx ids (#7431) * ROCm/AMD CI coverage: arch-table parity, native-Linux lib prepend, RDNA4 grouped_mm, discovery-based shell suite Three merged ROCm fixes shipped without tests, and the CI wiring that would have run them was gated on files the fixes do not touch. Tests added (113): tests/studio/install/test_rocm_arch_table_parity.py (27) diffs the four duplicated gfx -> AMD pip-index tables across install.sh, install.ps1, studio/setup.ps1 and install_python_stack.py, plus the GPU-name -> arch tables and the torch 2.11 pin allowlist. tests/studio/install/test_rocm_native_linux_lib_dirs.py (26) covers #7233: system-ROCm lib dirs prepended ahead of bundled libggml-hip, the /dev/kfd + not-WSL + libhsa gate, the opt-out env var, root resolution order, and source parity between the two copies. studio/backend/tests/test_grouped_mm_rdna4_fallback.py (46) covers #7292: registration on the CUDA dispatch key, grouped and ungrouped numerics, bias/dtype promotion, and the Linux HIP<7.13 + RDNA4 name gate, executed from the shipped source rather than a copy. tests/studio/test_ci_shell_suite_coverage.py (14) fails if either shell runner goes back to a hardcoded list or skips a file without a recorded reason. CI wiring: studio-backend-ci.yml: add install.sh / install.ps1 to the path filter (the suites it runs assert against those two files, so install-only changes -- the shape most AMD/ROCm routing fixes take -- skipped it), and replace the 13-file hardcoded shell list with directory discovery. That list had fallen seven files behind, including test_strixhalo_wsl_reroute.sh, the only shell coverage of the ROCm WSL reroute, which had never run on a PR. tests/run_all.sh: same discovery loop so local and CI agree. * Test review fixes: assert on outcomes, not on the code under test Self-review of the previous commit found four tests that passed for the wrong reason. 1. The arch-table parity test pinned expected gfx ids copied out of the shipped tables, which enshrined three upstream inaccuracies as correct: RX 9070 (non-XT) is gfx1201 not gfx1200, RX 7800 XT is gfx1101 not gfx1100, and PRO V710 is gfx1101 not gfx1102 per AMD's ROCm compatibility matrix. The expectation is now the AMD pip index leaf -- the thing the tables exist to produce, and what a wrong answer costs the user. The three known drifts are listed explicitly with a test asserting they stay cosmetic, i.e. that the wrong and right ids still map to the same wheel index. That test turns red the day one of them starts routing users to the wrong wheel. 2. The RDNA4 device-name test extracted the regex from worker.py and then matched with it, so it could not fail. Widening the pattern -- the dangerous edit, since it forces the slow Python mm fallback onto RDNA3 users -- would have been silently accepted. It now reads the live pattern and checks it against fixed cases, plus asserts the name match stays guarded by `not _lin_arch` and that the name is lowercased before matching. 3. The CI-coverage test matched a verbatim line of studio-backend-ci.yml, so reindenting the step would fail the build while a real regression to a hardcoded list could slip past a reformat. It now parses the YAML, finds the step by name, and asserts on the glob plus the absence of individual filenames. The path-filter test likewise reads the parsed trigger instead of scanning raw text. 4. A set comprehension in the parity helper had a ternary whose branches were identical. Mutation-tested: widening the RDNA4 regex, desyncing one copy of the name table, dropping install.sh from the path filter, and re-skipping the ROCm WSL shell suite each fail at least two tests. Verified on Linux (WSL Ubuntu 24.04) with CI's torch pin: 86 + 48 pass. * Fix three wrong gfx ids in the GPU-name arch tables The name -> gfx tables disagreed with AMD's ROCm compatibility matrix on three entries. Corrected against the "Radeon GPU" list at rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html: RX 9070, RX 9070 GRE gfx1200 -> gfx1201 (Navi 48, same die as the XT) RX 7800 XT, RX 7700 XT gfx1100 -> gfx1101 (Navi 32, not Navi 31) PRO W7700 gfx1100 -> gfx1101 PRO V710 gfx1102 -> gfx1101 (Navi 32, not Navi 33) No wheel changes for anyone: gfx1200/gfx1201 both resolve to gfx120X-all and gfx1100/gfx1101/gfx1102 all resolve to gfx110X-all, in all four copies of the index-family map. That collapse is why the errors survived being copied into six places -- the leaf-level tests could not see them. It was not purely cosmetic, though. install.sh's second copy feeds "Tip: set UNSLOTH_ROCM_GFX_ARCH=", so a 7800 XT user following the printed advice exported gfx1100 and made a wrong id authoritative for every later run. It would also have become a real misroute the moment AMD split a family across index leaves, as they already do for gfx1151/gfx1150. Fixed in all six places, which is two more than the table's own "kept in sync with" comments claim exist: install.sh _infer_amd_gfx_arch_from_gpu_name install.sh case "$_gpu_disp_mkt" (banner + env tip; undocumented) studio/setup.sh install.ps1 studio/setup.ps1 studio/install_python_stack.py Ordering is preserved: the gfx1102 arm still precedes gfx1101 in the shell copies so "RX 7700S" cannot fall onto the "RX 7700" glob, and the PowerShell copies keep the (?!S) lookahead. Test changes: - test_rocm_arch_table_parity.py gains _AMD_DOCUMENTED_ARCH, exact gfx ids transcribed from AMD rather than from the tables. Agreement between six copies proves nothing when all six were transcribed from the same mistake, so the ground truth has to come from outside. Verified it catches the bug: against the pre-fix tables it fails 6 tests. - The parity check now covers all six copies. It had four; the two install.sh copies were being treated as one, and _WIN_GPU_NAME_ARCH_TABLE was not checked at all. - test_rocm_support.py's TestGfxArchNameFallback pinned two of the wrong ids as expected values; updated, and extended with a 9060 XT and a 7900 XTX case so each RDNA3/4 die is represented. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard against unregistered copies of the GPU-name arch table Counting the copies by hand is what let them drift: the in-code "kept in sync with" comments claimed four, the arch-id fix found six, and scanning the tree turns up a seventh. TestNoUnregisteredArchTable rediscovers the copies from the source tree instead of trusting a hand-maintained list. A table line is one that names a card and gives its arch; real tables score 9-17 such lines and the only other hits in the repo are two single-line prose comments, so the three-line threshold is not load-bearing. A companion test asserts the scan still finds the known copies, so the heuristic cannot go blind and pass by finding nothing. The seventh copy is tests/_zoo_rocm_spoof.py, the fixture other ROCm tests build their fake AMD host from. It states the mapping backwards (gfx -> the name torch should report), which makes it an independent witness: it had gfx1101 -> RX 7800 XT and gfx1201 -> RX 9070 XT right while all six installer copies were wrong, and nothing compared the two. Now they are round-tripped against each other. RX 6700 XT is pinned as a known divergence rather than normalised. AMD's compatibility matrix documents no consumer RX 6000 card and no gfx1031 at all, the installer arm is commented "gfx103X family", and gfx1031 appears only as an index-family key, never as a value a name table emits. With no external source to correct against, changing shipped behaviour would be guesswork. A test fails if the divergence ever disappears, so the exemption cannot go stale. Also adds the reverse of the AMD-matrix check: a documented card that matches no arm anywhere is a silent CPU fallback rather than a wrong id. This cannot detect hardware nobody transcribed, which would need a live fetch of AMD's matrix and a non-hermetic suite; the docstring says so rather than implying coverage that is not there. Verified on Linux: 478 passed, plus all five new guards mutation-tested to confirm each fails when its invariant is broken. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Docstring said six copies; the list under it now has seven * tests: run discovered shell tests with bash, not sh tests/run_all.sh discovered tests/sh/ instead of listing files, but still invoked each one with sh. Every file there declares a bash shebang, and on Debian/Ubuntu /bin/sh is dash: test_apt_distro_prompt.sh, test_studio_home_node_dir.sh and test_with_llama_cpp_dir_link_behavior.sh fail on bashisms under dash and pass under bash. The old hand-written list happened to name only dash-clean files, so switching to discovery is what surfaced it. Backend CI already used bash, so this was a local-only break. Guarded by a new test asserting both runners invoke tests/sh/ with bash. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Krackan Point (Radeon 860M/840M) routed to the gfx1150 wheel index The GPU-name tables map 860M/840M and the Ryzen AI 7 350 / AI 5 340 CPU strings to gfx1150, but Krackan Point is gfx1152. AMD's own lemonade table (src/cpp/server/system_info.cpp) maps both Krackan iGPUs to gfx1152. Unlike the three ids already fixed here, this one is not wheel-neutral: repo.amd.com publishes gfx1150 and gfx1152 as separate index leaves with separately built torch wheels, so these laptops were installing wheels built for a different LLVM target. gfx1152 was absent from the codebase entirely, so it needed the index-family maps, the torch 2.11 floor lists (same _grouped_mm bug as gfx1150/1151), the Strix reroute set and the Windows arch allowlist as well as the seven name tables. The parity test added in this PR did not catch it because its AMD-matrix expectations stopped at 890M/880M. Added the APU rows, so the case that actually changes a wheel is now covered: reverting the tables fails 9 tests naming 860M, 840M and Krackan. gfx1153 (Ryzen AI 5 430 era) is left alone; AMD publishes no gfx1153 wheel family, so there is nothing to route it to. Verified: bash -n on both shell installers, PowerShell AST parse on both .ps1 files, python ast.parse on all touched modules, install suite 1334 passed with no new failures against main, shell suite 20 files. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add gfx1152 to unified-memory classifiers, make parity allowlist set-based Krackan Point (gfx1152, Radeon 860M/840M) is the third RDNA 3.5 APU and shares one GPU/system-RAM pool exactly like Strix Point (gfx1150) and Strix Halo (gfx1151), but only the installers knew about it. The two runtime classifiers still had two-element arch sets, so a Krackan laptop got the 0.90 discrete headroom factor on a shared pool and ran llama.cpp without GGML_CUDA_ENABLE_UNIFIED_MEMORY. - worker.py _rocm_classify_unified_memory: add gfx1152 to the arch set, and 860m/840m to the device-name fallback. The NVIDIA GeForce 840M cannot collide there: the function is only reached under _hw.IS_ROCM. - llama_cpp.py _amd_apu_wants_unified_memory: add gfx1152 to the arch set. - Tests for both, including the :sramecc-:xnack- suffix form. TestGfx211AllowlistParity compared four hardcoded allowlist strings, so adding gfx1152 to all four installers correctly turned three assertions red without any installer actually disagreeing with another. Each test now extracts the set its installer holds and compares it to one EXPECTED constant. Order and spacing are free, membership is not, and the next leaf is a one-line edit instead of four. * [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/studio-backend-ci.yml | 46 +- install.ps1 | 18 +- install.sh | 37 +- studio/backend/core/inference/llama_cpp.py | 6 +- studio/backend/core/training/worker.py | 17 +- .../tests/test_amd_apu_unified_memory.py | 4 +- .../tests/test_grouped_mm_rdna4_fallback.py | 418 +++++++++++ studio/backend/tests/test_rocm_oom_guard.py | 7 + studio/install_python_stack.py | 38 +- studio/setup.ps1 | 22 +- studio/setup.sh | 12 +- tests/_zoo_rocm_spoof.py | 3 +- tests/python/test_cross_platform_parity.py | 68 +- tests/run_all.sh | 33 +- .../install/test_rocm_arch_table_parity.py | 651 ++++++++++++++++++ .../test_rocm_native_linux_lib_dirs.py | 438 ++++++++++++ tests/studio/install/test_rocm_support.py | 25 +- tests/studio/test_ci_shell_suite_coverage.py | 197 ++++++ 18 files changed, 1923 insertions(+), 117 deletions(-) create mode 100644 studio/backend/tests/test_grouped_mm_rdna4_fallback.py create mode 100644 tests/studio/install/test_rocm_arch_table_parity.py create mode 100644 tests/studio/install/test_rocm_native_linux_lib_dirs.py create mode 100644 tests/studio/test_ci_shell_suite_coverage.py diff --git a/.github/workflows/studio-backend-ci.yml b/.github/workflows/studio-backend-ci.yml index 3968f2e80a..ec437e0c32 100644 --- a/.github/workflows/studio-backend-ci.yml +++ b/.github/workflows/studio-backend-ci.yml @@ -30,6 +30,13 @@ on: - 'unsloth/**' - 'unsloth_cli/**' - 'tests/**' + # The root installers: tests/sh/*.sh and tests/studio/install/* assert + # against these two files, so a change here must run the suite that + # covers it. Without them an install-only edit (the shape most AMD/ROCm + # routing fixes take) skipped Backend CI entirely. + - 'install.sh' + - 'install.ps1' + - 'scripts/**' - 'pyproject.toml' - '.github/workflows/studio-backend-ci.yml' push: @@ -217,27 +224,32 @@ jobs: tests/studio/test_xpu_spoof_pipeline.py - name: Shell installer tests - # Subset that does not depend on a writable / pristine install.sh - # tree; test_install_host_defaults.sh checks install.ps1 layout - # which has drifted (separate followup). + # Auto-discovered rather than allowlisted. The old hardcoded list had + # silently fallen seven files behind tests/run_all.sh, including + # test_strixhalo_wsl_reroute.sh -- the only shell coverage of the ROCm + # WSL reroute -- so that suite never ran on a PR. Skips are explicit, + # each with a reason, and tests/studio/test_ci_shell_suite_coverage.py + # fails if this step stops discovering the directory or the skip list + # grows without one. + # + # Skipped: + # test_install_host_defaults.sh: asserts an install.ps1 layout that + # has drifted (separate followup). + # test_install_rollback_lifecycle.sh: already runs on both platforms + # in cross-platform-parity-ci.yml. run: | set -e - for s in \ - tests/sh/test_get_torch_index_url.sh \ - tests/sh/test_mac_intel_compat.sh \ - tests/sh/test_node_decision.sh \ - tests/sh/test_studio_home_node_dir.sh \ - tests/sh/test_system_node_readonly.sh \ - tests/sh/test_nvcc_meets_llama_minimum.sh \ - tests/sh/test_resolve_cuda_archs.sh \ - tests/sh/test_staged_validation_enabled.sh \ - tests/sh/test_tauri_install_exit_order.sh \ - tests/sh/test_torch_constraint.sh \ - tests/sh/test_torch_flavor.sh \ - tests/sh/test_with_llama_cpp_dir_flag.sh \ - tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do + skip="test_install_host_defaults.sh test_install_rollback_lifecycle.sh" + found=0 + for s in tests/sh/test_*.sh; do + case " $skip " in + *" $(basename "$s") "*) echo "skipping $s (see workflow comment)"; continue ;; + esac + found=$((found + 1)) echo "::group::$s" bash "$s" echo "::endgroup::" done + [ "$found" -gt 0 ] || { echo "::error::no shell tests discovered under tests/sh"; exit 1; } + echo "ran $found shell installer test files" diff --git a/install.ps1 b/install.ps1 index 9c91d4ba16..a2aff0b69a 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1917,12 +1917,14 @@ exit 0 # (gfx120X/110X/1151/1150/103X); unknown names fall back to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( - @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080) - @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060) + @{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: RX 9070 XT / 9070 GRE / 9070 / 9080) + @{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: RX 9060 XT / 9060) @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) - @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) - @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) + @{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + @{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + @{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31) + @{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32) + @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33) @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family @@ -2203,6 +2205,7 @@ exit 0 $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) + "gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point) "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000) @@ -2224,6 +2227,7 @@ exit 0 $torchFloorMap = @{ "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" + "gfx1152" = "torch>=2.11.0,<2.12.0" } # Companion ranges track the torch ceiling so pip resolves a consistent # trio on AMD's per-arch index (each published independently). Mirrors @@ -2231,10 +2235,12 @@ exit 0 $torchvisionFloorMap = @{ "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" + "gfx1152" = "torchvision>=0.26.0,<0.27.0" } $torchaudioFloorMap = @{ "gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0" "gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0" + "gfx1152" = "torchaudio>=2.11.0,<2.12.0" } $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } if ($archFamily) { @@ -2264,7 +2270,7 @@ exit 0 $_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2) } # Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare. - $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf + $_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $_pinLeaf if ($_pinGfx211 -or $_pinRocm211) { $ROCmIndexUrl = $TorchIndexUrl $ROCmTorchFloor = "torch>=2.11.0,<2.12.0" diff --git a/install.sh b/install.sh index dface28918..3bc2ff4c88 100755 --- a/install.sh +++ b/install.sh @@ -2260,6 +2260,7 @@ _amd_arch_index_family_for_gfx() { gfx1201|gfx1200) echo gfx120X-all ;; gfx1151) echo gfx1151 ;; gfx1150) echo gfx1150 ;; + gfx1152) echo gfx1152 ;; gfx1103|gfx1102|gfx1101|gfx1100) echo gfx110X-all ;; gfx1036|gfx1035|gfx1034|gfx1033|gfx1032|gfx1031|gfx1030) echo gfx103X-all ;; gfx90a) echo gfx90a ;; @@ -2271,12 +2272,14 @@ _amd_arch_index_family_for_gfx() { # Map a GPU marketing name to gfx arch (kept in sync with install.ps1 nameArchTable). _infer_amd_gfx_arch_from_gpu_name() { case "$1" in - *"9070 XT"*|*9080*) echo gfx1201 ;; - *9070*|*9060*) echo gfx1200 ;; + *9070*|*9080*) echo gfx1201 ;; + *9060*) echo gfx1200 ;; *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) echo gfx1151 ;; - *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1150 ;; - *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) echo gfx1102 ;; - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) echo gfx1100 ;; + *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) echo gfx1150 ;; + *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) echo gfx1152 ;; + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) echo gfx1102 ;; + *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) echo gfx1101 ;; + *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) echo gfx1100 ;; *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) echo gfx1103 ;; *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) echo gfx1030 ;; *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) echo gfx1032 ;; @@ -2316,10 +2319,14 @@ _infer_linux_amd_gfx_arch() { echo gfx1151 return 0 fi - if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then + if [ -n "$_gpu_evidence" ] && grep -qiE '890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]' /proc/cpuinfo 2>/dev/null; then echo gfx1150 return 0 fi + if [ -n "$_gpu_evidence" ] && grep -qiE '860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33' /proc/cpuinfo 2>/dev/null; then + echo gfx1152 + return 0 + fi if command -v lspci >/dev/null 2>&1; then # A non-AMD controller can enumerate first (Intel/ASPEED before an AMD # dGPU), so scan every display-class line and take the first AMD one @@ -3055,7 +3062,7 @@ if [ "$_torch_index_pinned" = false ] && [ "$SKIP_TORCH" = false ] && \ # whole handoff (a user-set override re-exports unchanged). export UNSLOTH_ROCM_GFX_ARCH="$_linux_inferred_gfx" case "$_linux_inferred_gfx" in - gfx1201|gfx1200|gfx1151|gfx1150) + gfx1201|gfx1200|gfx1151|gfx1150|gfx1152) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" @@ -3124,7 +3131,7 @@ fi # and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a # custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced. case "$_torch_index_leaf" in - rocm7.2|gfx120x-all|gfx1151|gfx1150) + rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0" TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0" @@ -3243,7 +3250,7 @@ case "$_torch_index_leaf" in fi _strix_gfx="" case "$_runtime_gfx" in - gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;; + gfx1151|gfx1150|gfx1152) _strix_gfx="$_runtime_gfx" ;; esac # Skip rocm7.13+ generic indexes: they already ship the fixes, so the # arch build (rocm7.13) would be a downgrade rather than a rescue. @@ -3339,12 +3346,14 @@ elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then # gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on # gfx1102 (bash case has no negative lookahead like the PS tables). case "$_gpu_disp_mkt" in - *"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 - *9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 + *9070*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4 (Navi 48) + *9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4 (Navi 44) *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) - *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) + *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _gpu_disp_gfx="gfx1152" ;; # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3 (Navi 33) + *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) _gpu_disp_gfx="gfx1101" ;; # RDNA 3 (Navi 32) + *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point) *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _gpu_disp_gfx="gfx1030" ;; # RDNA 2 (Navi 21) *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _gpu_disp_gfx="gfx1032" ;; # RDNA 2 (Navi 23) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index d46e5a0fe1..e54b5269c1 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -3183,7 +3183,7 @@ class LlamaCppBackend: @staticmethod def _amd_apu_wants_unified_memory(gpu_indices = None) -> bool: - """True only for AMD unified-memory APUs (gfx1150/gfx1151), where + """True only for AMD unified-memory APUs (gfx1150/gfx1151/gfx1152), where GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM (it hurts discrete GPUs). gpu_indices (PHYSICAL ids) scopes the check to the selected GPUs, so a dGPU on a mixed host is not treated as unified-memory; @@ -3213,7 +3213,9 @@ class LlamaCppBackend: ) arch_by_id[pid] = _arch.split(":")[0].strip().lower() for _i in list(gpu_indices) if gpu_indices is not None else list(arch_by_id): - if arch_by_id.get(_i) in {"gfx1150", "gfx1151"}: + # gfx1152 is Krackan Point (Radeon 860M/840M), the third RDNA 3.5 + # APU: same shared GPU/system-RAM pool as Strix Point/Halo. + if arch_by_id.get(_i) in {"gfx1150", "gfx1151", "gfx1152"}: return True except Exception: return False diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 03327d3320..baf6329dae 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -764,8 +764,8 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: - ``gcn_arch``: canonical arch string (e.g. ``"gfx1151"``) when a known attribute is present, else ``""``. - ``is_unified``: ``True`` for AMD APUs with a shared GPU/system-RAM pool - (gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower - ``set_per_process_memory_fraction`` cap to leave OS headroom. + (gfx1150 Strix Point, gfx1151 Strix Halo, gfx1152 Krackan Point) — these + need a lower ``set_per_process_memory_fraction`` cap to leave OS headroom. Classification priority: 1. ``props.is_integrated`` truthy (hipDeviceProp_t.integrated -- the @@ -778,6 +778,7 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: - gfx1151 Strix Halo / Gorgon Halo: ``Radeon 8065S`` (Ryzen AI Max+ 495), ``Radeon 8060S`` (Ryzen AI MAX+ 395), ``Radeon 8050S`` (cut-down SKU) + - gfx1152 Krackan Point: ``Radeon 860M``, ``Radeon 840M`` """ gcn_arch = "" for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"): @@ -797,9 +798,13 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: return gcn_arch, True if gcn_arch: - return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"} + # gfx1152 is Krackan Point, the third RDNA 3.5 APU: same shared + # GPU/system-RAM pool as Strix Point (gfx1150) and Strix Halo (gfx1151). + return gcn_arch, gcn_arch in {"gfx1150", "gfx1151", "gfx1152"} - # Arch attrs absent — fall back to device-name matching. + # Arch attrs absent — fall back to device-name matching. Only reached under + # _hw.IS_ROCM, so the NVIDIA GeForce 840M cannot collide with the Krackan + # markers here. dev_lower = (getattr(props, "name", "") or "").lower() is_unified = ( "890m" in dev_lower @@ -807,6 +812,8 @@ def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]: or "8065s" in dev_lower or "8060s" in dev_lower or "8050s" in dev_lower + or "860m" in dev_lower + or "840m" in dev_lower ) return gcn_arch, is_unified @@ -2828,7 +2835,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) -> # On ROCm, exhausting VRAM can hang the HIP driver instead of raising. # set_per_process_memory_fraction caps the allocator so PyTorch raises # OutOfMemoryError first (NVIDIA already has a graceful OOM path). - # Unified-memory APUs (gfx1150/gfx1151) share GPU+system RAM, so use 0.80 + # Unified-memory APUs (gfx1150/gfx1151/gfx1152) share GPU+system RAM, so use 0.80 # vs 0.90 for discrete. Classify via gcnArchName, else device-name markers. # Non-fatal: skipped if torch is not importable. if _hw.IS_ROCM: diff --git a/studio/backend/tests/test_amd_apu_unified_memory.py b/studio/backend/tests/test_amd_apu_unified_memory.py index 9fd8260bf2..be85fd56d1 100644 --- a/studio/backend/tests/test_amd_apu_unified_memory.py +++ b/studio/backend/tests/test_amd_apu_unified_memory.py @@ -2,7 +2,7 @@ # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 """GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs -(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS.""" +(gfx1150/gfx1151/gfx1152), never for discrete AMD, NVIDIA, CPU or macOS.""" from __future__ import annotations @@ -35,6 +35,8 @@ def _fake_torch( [ ("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped) ("6.2.0", ["gfx1150"], True), # Strix Point APU + ("6.2.0", ["gfx1152"], True), # Krackan Point APU (Radeon 860M/840M) + ("6.2.0", ["gfx1152:sramecc-:xnack-"], True), # same, feature flags stripped ("6.2.0", ["gfx1100"], False), # discrete RDNA3 ("6.2.0", ["gfx1201"], False), # discrete RDNA4 ("6.2.0", ["gfx942"], False), # MI300X (data center) diff --git a/studio/backend/tests/test_grouped_mm_rdna4_fallback.py b/studio/backend/tests/test_grouped_mm_rdna4_fallback.py new file mode 100644 index 0000000000..675b9c3210 --- /dev/null +++ b/studio/backend/tests/test_grouped_mm_rdna4_fallback.py @@ -0,0 +1,418 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Numerics + gating for the RDNA4 _grouped_mm CPU fallback (PRs #7276 / #7292). + +RDNA4 (gfx1200/gfx1201) ships a null HIP `_grouped_mm` kernel on ROCm <= 7.12 +(fixed in 7.13, ROCm/TheRock #5284). Training MoE models there crashes with +0xC0000005 on Windows and a plain segfault on Linux, so worker.py registers a +Python mm/bmm fallback on the CUDA dispatch key. + +The fallback is silent, GPU-gated, and reimplements a matmul: if it is wrong, an +RX 9070 user does not crash, they train on quietly wrong gradients. Until now the +only coverage was `assert '_gm_lib.impl("_grouped_mm"' in source` -- the math was +never executed once, in any suite. + +worker.py cannot be imported here (module-level structlog/backend imports), so +`_install_grouped_mm_cpu_fallback` is lifted out with ast and driven with a fake +`torch_mod` that forwards to real CPU torch. That also pins the op surface: the +fallback may only use the ops the fake exposes, and the registration is captured +instead of hitting a real CUDA dispatch key that CI runners do not have. + +The two gates around it are exec'd straight out of the source so this file tests +the shipped expressions rather than a copy of them. +""" + +import ast +import re +import textwrap +from pathlib import Path +from types import SimpleNamespace + +import pytest +import torch + + +_WORKER_PATH = Path(__file__).resolve().parents[1] / "core" / "training" / "worker.py" +_WORKER_SOURCE = _WORKER_PATH.read_text(encoding = "utf-8") + + +def _load_installer(): + """exec just _install_grouped_mm_cpu_fallback out of worker.py.""" + tree = ast.parse(_WORKER_SOURCE) + fn = [ + n + for n in tree.body + if isinstance(n, ast.FunctionDef) and n.name == "_install_grouped_mm_cpu_fallback" + ] + assert fn, "_install_grouped_mm_cpu_fallback not found in core/training/worker.py" + ns: dict = {} + exec(compile(ast.Module(body = fn, type_ignores = []), str(_WORKER_PATH), "exec"), ns) + return ns["_install_grouped_mm_cpu_fallback"] + + +_install_grouped_mm_cpu_fallback = _load_installer() + + +class _RecordingLibrary: + """Stands in for torch.library.Library: captures the registration instead of + binding it to a CUDA dispatch key no CI runner has.""" + + def __init__(self, namespace, kind): + self.namespace = namespace + self.kind = kind + self.registrations = [] + + def impl(self, name, fn, dispatch_key): + self.registrations.append((name, fn, dispatch_key)) + + +class _RecordingLogger: + def __init__(self): + self.info_calls = [] + self.warning_calls = [] + + def info(self, *args, **kwargs): + self.info_calls.append(args) + + def warning(self, *args, **kwargs): + self.warning_calls.append(args) + + +def _fake_torch(): + """Real CPU torch behind the exact op surface the fallback is allowed to use. + + Anything else the fallback reaches for raises AttributeError here, which is + the point: a new dependency has to be a deliberate edit, not a silent one.""" + return SimpleNamespace( + library = SimpleNamespace(Library = _RecordingLibrary), + mm = torch.mm, + bmm = torch.bmm, + matmul = torch.matmul, + cat = torch.cat, + zeros = torch.zeros, + ) + + +@pytest.fixture +def fallback(): + """The registered _grouped_mm implementation, plus the Library it landed on.""" + torch_mod = _fake_torch() + logger = _RecordingLogger() + lib = _install_grouped_mm_cpu_fallback(torch_mod, logger, "test") + assert lib.registrations, "the fallback registered nothing" + name, fn, key = lib.registrations[0] + return SimpleNamespace(fn = fn, lib = lib, logger = logger, name = name, key = key) + + +class TestRegistration: + """Where the override lands. Getting the namespace or dispatch key wrong is a + silent no-op: training still crashes on the null HIP kernel.""" + + def test_overrides_aten_grouped_mm_on_the_cuda_key(self, fallback): + assert fallback.lib.namespace == "aten" + assert fallback.lib.kind == "IMPL" + assert fallback.name == "_grouped_mm" + # ROCm dispatches through the CUDA key; "HIP"/"PrivateUse1" would not bind. + assert fallback.key == "CUDA" + + def test_registers_exactly_once(self, fallback): + assert len(fallback.lib.registrations) == 1 + + def test_returns_the_library_so_the_caller_can_keep_it_alive(self, fallback): + """A dropped Library is garbage collected and the override silently + unregisters mid-run; worker.py parks it in a module global.""" + assert isinstance(fallback.lib, _RecordingLibrary) + assert "_WINDOWS_ROCM_GROUPED_MM_LIB = _install_grouped_mm_cpu_fallback(" in _WORKER_SOURCE + + def test_logs_the_patch_with_its_label(self, fallback): + assert fallback.logger.info_calls, "the patch must be visible in the run log" + assert "test" in fallback.logger.info_calls[0] + + +class TestUngroupedNumerics: + """offs=None: plain matmul, one path per rank combination. The 3-D case is + the regression #7292 fixed -- an unconditional mm() broke MoE experts.""" + + def test_2d_by_2d_matches_mm(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.mm(a, b)) + + def test_3d_by_3d_matches_bmm(self, fallback): + a = torch.randn(3, 6, 4) + b = torch.randn(3, 4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.bmm(a, b)) + + def test_3d_by_2d_matches_matmul(self, fallback): + a = torch.randn(3, 6, 4) + b = torch.randn(4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.matmul(a, b)) + + def test_2d_by_3d_matches_matmul(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(3, 4, 5) + torch.testing.assert_close(fallback.fn(a, b), torch.matmul(a, b)) + + def test_non_contiguous_inputs_are_handled(self, fallback): + """Transposed views reach _grouped_mm constantly; every path calls + .contiguous() and this catches it if one stops.""" + a = torch.randn(4, 6).t() + b = torch.randn(5, 4).t() + torch.testing.assert_close(fallback.fn(a, b), torch.mm(a, b)) + + +class TestGroupedNumerics: + """offs=[end-row of each group], the MoE token-routing layout.""" + + def test_matches_per_group_mm_with_3d_weights(self, fallback): + a = torch.randn(7, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([2, 5, 7]) + expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[1], a[5:7] @ b[2]], dim = 0) + torch.testing.assert_close(fallback.fn(a, b, offs), expected) + + def test_shared_2d_weight_is_reused_for_every_group(self, fallback): + a = torch.randn(7, 4) + b = torch.randn(4, 5) + offs = torch.tensor([2, 5, 7]) + torch.testing.assert_close(fallback.fn(a, b, offs), a @ b) + + def test_empty_group_produces_no_rows(self, fallback): + """An expert that routed zero tokens (offs[i] == offs[i-1]) must + contribute nothing, not a stray row.""" + a = torch.randn(5, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([2, 2, 5]) + expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[2]], dim = 0) + got = fallback.fn(a, b, offs) + assert got.shape == (5, 5) + torch.testing.assert_close(got, expected) + + def test_rows_past_the_last_offset_are_not_dropped(self, fallback): + """Trailing tokens beyond offs[-1] go through the last expert; dropping + them would silently shrink the output instead of raising.""" + a = torch.randn(7, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([2, 5]) + expected = torch.cat([a[0:2] @ b[0], a[2:5] @ b[1], a[5:7] @ b[-1]], dim = 0) + got = fallback.fn(a, b, offs) + assert got.shape[0] == a.shape[0] + torch.testing.assert_close(got, expected) + + def test_zero_rows_returns_an_empty_result_not_an_error(self, fallback): + a = torch.randn(0, 4) + b = torch.randn(3, 4, 5) + offs = torch.tensor([], dtype = torch.int64) + got = fallback.fn(a, b, offs) + assert got.shape == (0, 5) + assert got.dtype == a.dtype + + def test_offsets_may_arrive_as_a_device_tensor_of_any_int_dtype(self, fallback): + a = torch.randn(4, 4) + b = torch.randn(2, 4, 5) + expected = torch.cat([a[0:2] @ b[0], a[2:4] @ b[1]], dim = 0) + for dtype in (torch.int32, torch.int64): + torch.testing.assert_close( + fallback.fn(a, b, torch.tensor([2, 4], dtype = dtype)), expected + ) + + +class TestBiasAndDtype: + def test_bias_is_added(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(4, 5) + bias = torch.randn(5) + torch.testing.assert_close(fallback.fn(a, b, None, bias), torch.mm(a, b) + bias) + + def test_bias_is_added_on_the_grouped_path_too(self, fallback): + a = torch.randn(4, 4) + b = torch.randn(2, 4, 5) + bias = torch.randn(5) + offs = torch.tensor([2, 4]) + expected = torch.cat([a[0:2] @ b[0], a[2:4] @ b[1]], dim = 0) + bias + torch.testing.assert_close(fallback.fn(a, b, offs, bias), expected) + + def test_out_dtype_is_honoured(self, fallback): + a = torch.randn(6, 4) + b = torch.randn(4, 5) + got = fallback.fn(a, b, None, None, torch.float64) + assert got.dtype == torch.float64 + torch.testing.assert_close(got, torch.mm(a, b).to(torch.float64)) + + def test_promotion_from_bias_is_cast_back_to_the_input_dtype(self, fallback): + """Without the restore, a promoted result changes the autograd dtype + downstream of every MoE layer.""" + a = torch.randn(6, 4, dtype = torch.float32) + b = torch.randn(4, 5, dtype = torch.float32) + bias = torch.randn(5, dtype = torch.float64) + got = fallback.fn(a, b, None, bias) + assert got.dtype == torch.float32 + + def test_out_dtype_wins_over_the_input_dtype_restore(self, fallback): + a = torch.randn(6, 4, dtype = torch.float32) + b = torch.randn(4, 5, dtype = torch.float32) + bias = torch.randn(5, dtype = torch.float64) + got = fallback.fn(a, b, None, bias, torch.float64) + assert got.dtype == torch.float64 + + def test_bf16_inputs_stay_bf16(self, fallback): + """The dtype training actually runs in.""" + a = torch.randn(6, 4).to(torch.bfloat16) + b = torch.randn(4, 5).to(torch.bfloat16) + got = fallback.fn(a, b) + assert got.dtype == torch.bfloat16 + torch.testing.assert_close(got.float(), (a.float() @ b.float()), rtol = 2e-2, atol = 2e-2) + + +def _exec_source_snippet(anchor: str, last_line: str, **variables): + """Run a slice of worker.py verbatim, so the gate under test is the shipped + one and not a copy that can drift.""" + start = _WORKER_SOURCE.find(anchor) + assert start != -1, f"gate snippet not found in worker.py: {anchor!r}" + start = _WORKER_SOURCE.rfind("\n", 0, start) + 1 # keep the indent for dedent() + end = _WORKER_SOURCE.find(last_line, start) + assert end != -1, f"end of gate snippet not found: {last_line!r}" + snippet = textwrap.dedent(_WORKER_SOURCE[start : end + len(last_line)]) + ns = {"re": re, **variables} + exec(compile(snippet, str(_WORKER_PATH), "exec"), ns) + return ns + + +class TestLinuxHipVersionGate: + """PR #7292's Linux gate. Too low a floor keeps the slow Python fallback on + fixed ROCm 7.13+; too high reintroduces the segfault on 7.12.""" + + _ANCHOR = '_m = re.match(r"(\\d+)\\.(\\d+)", _hip_str)' + _LAST = '_hip_lt_713 = "rocmsdk" not in _ver' + + def _decide(self, hip_str, version): + ns = _exec_source_snippet(self._ANCHOR, self._LAST, _hip_str = hip_str, _ver = version.lower()) + return ns["_hip_lt_713"] + + @pytest.mark.parametrize( + "hip_str,version,affected", + [ + ("7.12.0", "2.10.0+rocm7.12.0", True), # the broken kernel + ("7.6.0", "2.9.0+rocm7.6.0", True), + ("6.4.0", "2.8.0+rocm6.4.0", True), + ("7.13.0", "2.11.0+rocm7.13.0", False), # AMD's fix + ("7.14.0", "2.11.0+rocm7.14.0", False), + ("8.0.0", "2.12.0+rocm8.0.0", False), + ], + ) + def test_torch_version_hip_decides_when_present(self, hip_str, version, affected): + assert self._decide(hip_str, version) is affected + + @pytest.mark.parametrize( + "version,affected", + [ + ("2.10.0+rocm7.12.0", True), + ("2.11.0+rocm7.13.0", False), + ("2.11.0+rocm7.14.0", False), + ], + ) + def test_falls_back_to_the_rocm_tag_in_torch_version(self, version, affected): + """AMD SDK / Radeon wheels leave torch.version.hip unset.""" + assert self._decide("", version) is affected + + def test_unknown_version_is_assumed_affected(self): + """Fallback is slow but correct; a missed guard is a crash.""" + assert self._decide("", "2.9.0+unknown") is True + + def test_rocmsdk_wheels_without_a_version_are_assumed_fixed(self): + """rocmsdk wheels post-date the gfx120X fix.""" + assert self._decide("", "2.10.0+rocmsdk20260107") is False + + +class TestLinuxRdna4NameMatch: + """The name regex is the fallback when a wheel omits gcnArchName.""" + + def _pattern(self): + """Read whatever pattern worker.py currently uses, not a copy of the one + it used when this test was written. Anchoring on the literal pattern text + would make a *widened* regex -- the dangerous edit, since it silently + forces the slow Python fallback onto RDNA3 users -- fail as "moved" + instead of being checked against the cases below.""" + m = re.search(r"re\.search\(r\"([^\"]+)\",\s*_lin_name\)", _WORKER_SOURCE) + assert m, "could not locate the RDNA4 device-name regex in worker.py" + return m.group(1) + + def test_name_is_lowercased_before_matching(self): + """The pattern is all-lowercase, so it only works against a lowercased + name. Device names arrive mixed case ("AMD Radeon RX 9070 XT").""" + assert self._pattern() == self._pattern().lower(), "pattern is not all-lowercase" + assert re.search( + r"_lin_name\s*=\s*\(getattr\(_props,\s*\"name\",\s*\"\"\)\s*or\s*\"\"\)\.lower\(\)", + _WORKER_SOURCE, + ), "worker.py must lowercase the device name before matching the RDNA4 pattern" + + def test_name_match_is_only_a_fallback_when_arch_is_unknown(self): + """gcnArchName is authoritative when present. Letting the name regex fire + alongside a known arch would misclassify any card whose marketing name + happens to look RDNA4.""" + assert re.search( + r"not _lin_arch and re\.search\(r\"[^\"]+\",\s*_lin_name\)", _WORKER_SOURCE + ), "the RDNA4 name regex must be guarded by `not _lin_arch`" + + @pytest.mark.parametrize( + "name,is_rdna4", + [ + ("AMD Radeon RX 9070 XT", True), + ("AMD Radeon RX 9060 XT", True), + ("Radeon RX9070", True), + ("AMD Radeon AI PRO R9700", True), + ("AMD Radeon RX 7900 XTX", False), # RDNA3, kernel is fine + ("AMD Radeon 8060S Graphics", False), # Strix Halo + ("AMD Radeon RX 6800 XT", False), + ("NVIDIA GeForce RTX 4090", False), + ], + ) + def test_matches_only_rdna4_cards(self, name, is_rdna4): + assert bool(re.search(self._pattern(), name.lower())) is is_rdna4 + + +class TestLinuxGateStructure: + """The block is a few hundred lines into run_training_process and can only be + checked structurally; these pin the parts a refactor would quietly drop.""" + + def _linux_block(self): + start = _WORKER_SOURCE.find("1f-linux") + assert start != -1, "the Linux ROCm gfx120X guard (#7292) is gone from worker.py" + end = _WORKER_SOURCE.find("1g.", start) + assert end != -1 + return _WORKER_SOURCE[start:end] + + def test_gated_on_linux_and_rocm(self): + block = self._linux_block() + assert 'sys.platform.startswith("linux")' in block + assert "_hw.IS_ROCM" in block, "guard must not run on NVIDIA/CPU hosts" + + def test_requires_both_rdna4_and_an_affected_hip(self): + block = self._linux_block() + assert "if _rdna4 and _hip_lt_713:" in block + + def test_scans_every_visible_device(self): + """device_map="balanced" can place layers on a later card, so checking + device 0 alone misses the RDNA4 GPU.""" + block = self._linux_block() + assert "for _i in range(_torch_lin.cuda.device_count()):" in block + + def test_matches_both_rdna4_arch_ids(self): + block = self._linux_block() + assert '("gfx1200", "gfx1201")' in block + + def test_failure_to_patch_is_non_fatal(self): + """A broken patch attempt must not take down the whole training run.""" + block = self._linux_block() + assert "except Exception" in block + assert "logger.warning" in block + + def test_windows_and_linux_share_one_implementation(self): + """Two copies of this fallback would drift; #7292 deliberately hoisted it.""" + assert _WORKER_SOURCE.count("def _install_grouped_mm_cpu_fallback(") == 1 + assert _WORKER_SOURCE.count("_install_grouped_mm_cpu_fallback(") >= 3 # def + win32 + linux + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/studio/backend/tests/test_rocm_oom_guard.py b/studio/backend/tests/test_rocm_oom_guard.py index ad46f6ee41..5cdbe4f2a5 100644 --- a/studio/backend/tests/test_rocm_oom_guard.py +++ b/studio/backend/tests/test_rocm_oom_guard.py @@ -80,6 +80,7 @@ class TestCanonicalGcnArchName: [ ("gfx1150", True), # Strix Point ("gfx1151", True), # Strix Halo + ("gfx1152", True), # Krackan Point (Radeon 860M/840M) ("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete ("gfx906", False), # MI50 — discrete server GPU ("gfx1201", False), # RX 9070 XT — discrete @@ -166,9 +167,15 @@ class TestDeviceNameFallback: # gfx1151 Gorgon Halo (Ryzen AI Max 400 refresh) "Radeon 8065S Graphics", # Ryzen AI Max+ 495 "AMD Radeon 8065S", + # gfx1152 Krackan Point (Ryzen AI 7 350 / AI 5 340) + "Radeon 860M", + "AMD Radeon 860M Graphics", + "Radeon 840M", + "AMD Radeon 840M Graphics", # case variants "RADEON 8060S GRAPHICS", "radeon 8050s", + "RADEON 860M", ], ) def test_unified_memory_detected(self, device_name: str) -> None: diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index a29ba0d7e5..8c33cb6ce9 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -96,7 +96,9 @@ def _strix_needs_amd_arch_index(ver: tuple[int, int]) -> bool: # AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). # Mirrors *FloorMap in install.ps1 / setup.ps1; other arches ship <2.11 and stay bare. -_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset({"gfx120x-all", "gfx1151", "gfx1150"}) +_ROCM_GFX_TORCH211_LEAVES: frozenset[str] = frozenset( + {"gfx120x-all", "gfx1151", "gfx1150", "gfx1152"} +) # pytorch.org rocmX.Y indexes KNOWN to ship torch 2.11 (rocm7.2 only today); don't # floor an unknown newer rocm speculatively. Match install.sh / setup.ps1 / install.ps1. @@ -124,6 +126,7 @@ _WINDOWS_ROCM_TORCH_PKG_SPECS: dict[str, tuple[str, str, str]] = { "gfx1200": _ROCM_TORCH_PKG_SPECS["rocm7.2"], "gfx1151": _ROCM_TORCH_PKG_SPECS["rocm7.2"], "gfx1150": _ROCM_TORCH_PKG_SPECS["rocm7.2"], + "gfx1152": _ROCM_TORCH_PKG_SPECS["rocm7.2"], } _PYTORCH_WHL_BASE = ( os.environ.get("UNSLOTH_PYTORCH_MIRROR") or "https://download.pytorch.org/whl" @@ -369,6 +372,7 @@ _GFX_TO_AMD_INDEX_ARCH: dict[str, str] = { "gfx1200": "gfx120X-all", # RDNA 4 "gfx1151": "gfx1151", "gfx1150": "gfx1150", # RDNA 3.5 (Strix Halo/Point) + "gfx1152": "gfx1152", # RDNA 3.5 (Krackan Point) "gfx1103": "gfx110X-all", "gfx1102": "gfx110X-all", # RDNA 3 "gfx1101": "gfx110X-all", @@ -738,19 +742,18 @@ def _detect_windows_gfx_arch() -> str | None: # prebuilts / AMD Windows torch indexes support; unknown names return None # (callers then fall back cleanly to CPU). _WIN_GPU_NAME_ARCH_TABLE: "list[tuple[str, str]]" = [ - (r"9070 XT|9080", "gfx1201"), # RDNA 4 (Radeon RX 9070 XT / 9080) - (r"9070|9060", "gfx1200"), # RDNA 4 (Radeon RX 9070 / 9060) + (r"9070|9080", "gfx1201"), # RDNA 4 (Navi 48: Radeon RX 9070 XT / 9070 GRE / 9070 / 9080) + (r"9060", "gfx1200"), # RDNA 4 (Navi 44: Radeon RX 9060 XT / 9060) # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) (r"8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max", "gfx1151"), - # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - ( - r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]" - r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", - "gfx1150", - ), + # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + (r"890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]", "gfx1150"), + # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + (r"860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", "gfx1152"), # RDNA 3 desktop / workstation (Navi 31) - (r"RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700", "gfx1100"), - (r"RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710", "gfx1102"), # Navi 33 + (r"RX 7900|PRO W7900|PRO W7800", "gfx1100"), + (r"RX 7800|RX 7700(?!S)|PRO W7700|PRO V710", "gfx1101"), # Navi 32 + (r"RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500", "gfx1102"), # Navi 33 # RDNA 3 iGPU (Phoenix / Hawk Point) (r"780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme", "gfx1103"), (r"RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900", "gfx1030"), # Navi 21 @@ -777,13 +780,12 @@ def _linux_amd_gfx_from_cpuinfo() -> "str | None": return None if re.search(r"Ryzen AI Max|Radeon 80[0-9][05]S|Strix Halo", text, re.IGNORECASE): return "gfx1151" - if re.search( - r"890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]" - r"|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", - text, - re.IGNORECASE, - ): + if re.search(r"890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]", text, re.IGNORECASE): return "gfx1150" + if re.search( + r"860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33", text, re.IGNORECASE + ): + return "gfx1152" return None @@ -1896,7 +1898,7 @@ def _ensure_rocm_torch() -> None: # An explicit ROCm pin is authoritative: never auto-reroute it. if _strix_needs_amd_arch_index(ver) and _explicit_rocm_torch_index_url() is None: gfx_codes = _detect_amd_gfx_codes() - _strix_gfx = {"gfx1151", "gfx1150"} + _strix_gfx = {"gfx1151", "gfx1150", "gfx1152"} _detected_strix = _strix_gfx.intersection(gfx_codes) if _detected_strix: # Runtime-visible GPU (HIP_VISIBLE_DEVICES index into gfx_codes, else first); diff --git a/studio/setup.ps1 b/studio/setup.ps1 index c6932121c9..6a8499b195 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -484,7 +484,7 @@ function Redact-InstallOutput { # the install-spec path below and the other installers; other leaves ship <2.11 and stay default. function Test-RocmGfx211Leaf { param([string]$Leaf) - return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf + return @('gfx120x-all', 'gfx1151', 'gfx1150', 'gfx1152') -contains $Leaf } # rocmX.Y versions KNOWN to ship torch 2.11: rocm7.2 only today. Do NOT floor an unknown newer @@ -1496,12 +1496,14 @@ if (-not $HasNvidiaSmi) { # (gfx120X/110X/1151/1150/103X); unknown names fall back cleanly to CPU. elseif ($ROCmGpuLabel) { $nameArchTable = @( - @{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (Radeon RX 9070 XT / 9080) - @{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (Radeon RX 9070 / 9060) + @{ P = "9070|9080"; A = "gfx1201" } # RDNA 4 (Navi 48: Radeon RX 9070 XT / 9070 GRE / 9070 / 9080) + @{ P = "9060"; A = "gfx1200" } # RDNA 4 (Navi 44: Radeon RX 9060 XT / 9060) @{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) - @{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - @{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop / workstation (Navi 31) - @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33) + @{ P = "890M|880M|Strix Point|HX 37[05]|AI 9 HX|AI 9 36[05]"; A = "gfx1150" } # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + @{ P = "860M|840M|Krackan|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1152" } # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + @{ P = "RX 7900|PRO W7900|PRO W7800"; A = "gfx1100" } # RDNA 3 desktop / workstation (Navi 31) + @{ P = "RX 7800|RX 7700(?!S)|PRO W7700|PRO V710"; A = "gfx1101" } # RDNA 3 (Navi 32) + @{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500"; A = "gfx1102" } # RDNA 3 (Navi 33) @{ P = "780M|760M|740M|Phoenix|Hawk Point|Z1 Extreme|Z2 Extreme"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix / Hawk Point) @{ P = "RX 6900|RX 6800|RX 6750|RX 6700|PRO W6800|PRO W6900"; A = "gfx1030" } # RDNA 2 (Navi 21) -- gfx103X family @{ P = "RX 6650|RX 6600|PRO W6600|PRO W6650"; A = "gfx1032" } # RDNA 2 (Navi 23) -- gfx103X family @@ -2773,7 +2775,7 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode # for those or a correct CPU venv rebuilds every update. $_rocmWheelArches = @( "gfx1201", "gfx1200", # RDNA 4 - "gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point) + "gfx1151", "gfx1150", "gfx1152", # RDNA 3.5 (Strix Halo/Point, Krackan Point) "gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3 "gfx1036", "gfx1035", "gfx1034", "gfx1033", "gfx1032", "gfx1031", "gfx1030", # RDNA 2 (RX 6000) "gfx90a", "gfx908" # MI200 / MI100 @@ -3064,6 +3066,7 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu $archFamilyMap = @{ "gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4 "gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point) + "gfx1152" = "gfx1152" # RDNA 3.5 (Krackan Point) "gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3 "gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all" "gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000) @@ -3078,6 +3081,7 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu $torchFloorMap = @{ "gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0" "gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0" + "gfx1152" = "torch>=2.11.0,<2.12.0" } # Companion ranges for torchvision/torchaudio -- must stay in sync with the # torch ceiling so pip can always find a consistent trio on AMD's per-arch @@ -3089,10 +3093,12 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu $torchvisionFloorMap = @{ "gfx1201" = "torchvision>=0.26.0,<0.27.0"; "gfx1200" = "torchvision>=0.26.0,<0.27.0" "gfx1151" = "torchvision>=0.26.0,<0.27.0"; "gfx1150" = "torchvision>=0.26.0,<0.27.0" + "gfx1152" = "torchvision>=0.26.0,<0.27.0" } $torchaudioFloorMap = @{ "gfx1201" = "torchaudio>=2.11.0,<2.12.0"; "gfx1200" = "torchaudio>=2.11.0,<2.12.0" "gfx1151" = "torchaudio>=2.11.0,<2.12.0"; "gfx1150" = "torchaudio>=2.11.0,<2.12.0" + "gfx1152" = "torchaudio>=2.11.0,<2.12.0" } $archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null } $ROCmTorchSpec = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { "torch" } @@ -3104,7 +3110,7 @@ if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu # GPU arch detected but not in the supported wheel map — warn explicitly # so the user knows why they are getting CPU PyTorch instead of ROCm. substep "[WARN] AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow" - substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx1030-1036 (RDNA 2), gfx90a, gfx908" "Yellow" + substep " Supported: gfx1200/1201 (RDNA 4), gfx1150/1151/1152 (RDNA 3.5), gfx1100-1103 (RDNA 3), gfx1030-1036 (RDNA 2), gfx90a, gfx908" "Yellow" } else { # HIP SDK present ($HasROCm=true via amd-smi) but gcnArchName was not # readable — warn rather than silently falling back to CPU PyTorch. diff --git a/studio/setup.sh b/studio/setup.sh index f6a6bc346b..37d8154e59 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1167,12 +1167,14 @@ elif [ "$_setup_amd_detected" = true ]; then # gfx1102 matched BEFORE gfx1100 so the spaceless "RX 7700S" lands on # gfx1102 (bash case has no negative lookahead like the PS tables). case "$_setup_mkt" in - *"9070 XT"*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 - *9070*|*9060*) _setup_gfx="gfx1200" ;; # RDNA 4 + *9070*|*9080*) _setup_gfx="gfx1201" ;; # RDNA 4 (Navi 48) + *9060*) _setup_gfx="gfx1200" ;; # RDNA 4 (Navi 44) *"8065S"*|*"8060S"*|*"8050S"*|*"8040S"*|*"Strix Halo"*|*"Ryzen AI Max"*|*"AI Max"*) _setup_gfx="gfx1151" ;; # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+) - *"890M"*|*"880M"*|*"860M"*|*"840M"*|*"Strix Point"*|*"Krackan"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375) - *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*|*"PRO V710"*) _setup_gfx="gfx1102" ;; # RDNA 3 (Navi 33) - *"RX 7900"*|*"RX 7800"*|*"RX 7700"*|*"PRO W7900"*|*"PRO W7800"*|*"PRO W7700"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) + *"890M"*|*"880M"*|*"Strix Point"*|*"HX 37"*|*"AI 9 HX"*|*"AI 9 36"*) _setup_gfx="gfx1150" ;; # RDNA 3.5 (Strix Point: Radeon 890M/880M, Ryzen AI 9 HX 370/375) + *"860M"*|*"840M"*|*"Krackan"*|*"AI 7 35"*|*"AI 5 34"*|*"AI 7 PRO 35"*|*"AI 5 33"*) _setup_gfx="gfx1152" ;; # RDNA 3.5 (Krackan Point: Radeon 860M/840M, Ryzen AI 7 350 / AI 5 340) + *"RX 7600"*|*"RX 7700S"*|*"RX 7650"*|*"PRO W7600"*|*"PRO W7500"*) _setup_gfx="gfx1102" ;; # RDNA 3 (Navi 33) + *"RX 7800"*|*"RX 7700"*|*"PRO W7700"*|*"PRO V710"*) _setup_gfx="gfx1101" ;; # RDNA 3 (Navi 32) + *"RX 7900"*|*"PRO W7900"*|*"PRO W7800"*) _setup_gfx="gfx1100" ;; # RDNA 3 desktop / workstation (Navi 31) *"780M"*|*"760M"*|*"740M"*|*"Phoenix"*|*"Hawk Point"*|*"Z1 Extreme"*|*"Z2 Extreme"*) _setup_gfx="gfx1103" ;; # RDNA 3 iGPU (Phoenix / Hawk Point) *"RX 6900"*|*"RX 6800"*|*"RX 6750"*|*"RX 6700"*|*"PRO W6800"*|*"PRO W6900"*) _setup_gfx="gfx1030" ;; # RDNA 2 (Navi 21) *"RX 6650"*|*"RX 6600"*|*"PRO W6600"*|*"PRO W6650"*) _setup_gfx="gfx1032" ;; # RDNA 2 (Navi 23) diff --git a/tests/_zoo_rocm_spoof.py b/tests/_zoo_rocm_spoof.py index 050191e9d1..f5dd090148 100644 --- a/tests/_zoo_rocm_spoof.py +++ b/tests/_zoo_rocm_spoof.py @@ -25,6 +25,7 @@ _PROFILES: dict[str, tuple[str, tuple[int, int], str]] = { "gfx1101": ("AMD Radeon RX 7800 XT", (11, 0), "6.4.43483"), "gfx1102": ("AMD Radeon RX 7600", (11, 0), "7.2.1"), "gfx1150": ("AMD Radeon 890M", (11, 5), "7.2.1"), # RDNA3.5 APU + "gfx1152": ("AMD Radeon 860M", (11, 5), "7.2.1"), "gfx1151": ("AMD Radeon 8060S", (11, 5), "7.2.1"), "gfx1200": ("AMD Radeon RX 9060 XT", (12, 0), "7.2.1"), # RDNA4 "gfx1201": ("AMD Radeon RX 9070 XT", (12, 0), "7.2.1"), @@ -73,7 +74,7 @@ def apply(gfx: str = "gfx1100", device_count: int = 1) -> None: _p.total_memory = 16 * 1024**3 _p.multi_processor_count = 40 _p.warp_size = 32 # RDNA wavefront (CDNA is 64) - _p.is_integrated = gfx in ("gfx1150", "gfx1151") + _p.is_integrated = gfx in ("gfx1150", "gfx1151", "gfx1152") _p.is_multi_gpu_board = False torch.cuda.get_device_properties = lambda *a, **k: _p diff --git a/tests/python/test_cross_platform_parity.py b/tests/python/test_cross_platform_parity.py index b3a9b99c55..b0a5c763d4 100644 --- a/tests/python/test_cross_platform_parity.py +++ b/tests/python/test_cross_platform_parity.py @@ -249,22 +249,42 @@ class TestTorchIndexOverrideParity: class TestGfx211AllowlistParity: - """The gfx per-arch 2.11-floor leaves (gfx120X-all / gfx1151 / gfx1150) must be the - SAME set in every installer and its stale/mismatch check. When they diverged, a - pinned gfx110X-all / gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update.""" + """The gfx per-arch 2.11-floor leaves must be the SAME set in every installer + and its stale/mismatch check. When they diverged, a pinned gfx110X-all / + gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update. - EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150"} + Each test extracts the set each installer actually holds and compares it + against EXPECTED, rather than matching one hardcoded ordering. Order and + spacing are free; membership is not. The earlier literal-string form had to + be edited in four places whenever a leaf was added, which is how adding + gfx1152 (Krackan Point) turned this class red without any installer + actually disagreeing with another.""" + + EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150", "gfx1152"} + + @staticmethod + def _leaves(blob: str) -> set[str]: + """The gfx leaves named in an allowlist literal, quoting-agnostic.""" + return set(re.findall(r"gfx[0-9a-z-]+", blob.lower())) def test_install_sh_allowlist(self): text = INSTALL_SH.read_text(encoding = "utf-8").lower() - # install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx120x-all|gfx1151|gfx1150). - m = re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150", text) + # install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx...|gfx...). + m = re.search(r"^\s*(rocm7\.2\|[a-z0-9|.\-]*)\)", text, re.MULTILINE) assert m, "install.sh gfx-2.11 allowlist case not found / changed" + assert self._leaves(m.group(1)) == self.EXPECTED, ( + f"install.sh gfx-2.11 allowlist is {sorted(self._leaves(m.group(1)))}, " + f"expected {sorted(self.EXPECTED)}" + ) def test_install_ps1_allowlist(self): text = INSTALL_PS1.read_text(encoding = "utf-8").lower() - m = re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text) + m = re.search(r"\$_pingfx211\s*=\s*@\(([^)]*)\)", text) assert m, "install.ps1 $_pinGfx211 allowlist not found / changed" + assert self._leaves(m.group(1)) == self.EXPECTED, ( + f"install.ps1 $_pinGfx211 is {sorted(self._leaves(m.group(1)))}, " + f"expected {sorted(self.EXPECTED)}" + ) def test_setup_ps1_defines_single_allowlist_helper(self): # setup.ps1 must define the allowlist once (Test-RocmGfx211Leaf) and reuse it, so @@ -273,9 +293,12 @@ class TestGfx211AllowlistParity: assert ( "function Test-RocmGfx211Leaf" in text ), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper" - assert re.search( - r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower() - ), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist" + m = re.search(r"function test-rocmgfx211leaf[\s\S]{0,400}?@\(([^)]*)\)", text.lower()) + assert m, "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist" + assert self._leaves(m.group(1)) == self.EXPECTED, ( + f"Test-RocmGfx211Leaf holds {sorted(self._leaves(m.group(1)))}, " + f"expected {sorted(self.EXPECTED)}" + ) assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, ( "setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not " "re-hardcode the allowlist (they must not diverge)" @@ -283,9 +306,12 @@ class TestGfx211AllowlistParity: def test_stack_py_allowlist(self): text = STACK_PY.read_text(encoding = "utf-8").lower() - assert ( - '"gfx120x-all", "gfx1151", "gfx1150"' in text - ), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed" + m = re.search(r"_rocm_gfx_torch211_leaves[^=]*=\s*frozenset\(\s*\{([^}]*)\}", text) + assert m, "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed" + assert self._leaves(m.group(1)) == self.EXPECTED, ( + f"_ROCM_GFX_TORCH211_LEAVES is {sorted(self._leaves(m.group(1)))}, " + f"expected {sorted(self.EXPECTED)}" + ) class TestCudaLeafDigitParity: @@ -351,15 +377,21 @@ class TestCudaLeafDigitParity: class TestKnown211SetParity: """The KNOWN-2.11 rocm/gfx set must be identical across all four installers: - exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}. + exactly {rocm7.2} plus TestGfx211AllowlistParity.EXPECTED. rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively.""" def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self): text = INSTALL_SH.read_text(encoding = "utf-8") - # The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves. - assert re.search( - r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text - ), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150" + # The 2.11 floor case matches exactly rocm7.2 + the gfx allowlist, in + # any order: it is the same set as TestGfx211AllowlistParity.EXPECTED, + # asserted here so the rocm-version half cannot drift on its own. + m = re.search(r"^\s*(rocm7\.2\|[a-zA-Z0-9|.\-]*)\)", text, re.MULTILINE) + assert m, "install.sh 2.11 floor case (rocm7.2|gfx...) not found / changed" + alternatives = set(m.group(1).lower().split("|")) + assert alternatives == {"rocm7.2"} | TestGfx211AllowlistParity.EXPECTED, ( + f"install.sh 2.11 floor is {sorted(alternatives)}, expected " + f"{sorted({'rocm7.2'} | TestGfx211AllowlistParity.EXPECTED)}" + ) # No speculative rocm7.3 anywhere. assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3" diff --git a/tests/run_all.sh b/tests/run_all.sh index 6eccffc75f..2497ba8fe7 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -7,18 +7,27 @@ set -e TESTS_DIR="$(cd "$(dirname "$0")" && pwd)" echo "=== Bash tests ===" -sh "$TESTS_DIR/sh/test_get_torch_index_url.sh" -sh "$TESTS_DIR/sh/test_mac_intel_compat.sh" -sh "$TESTS_DIR/sh/test_torch_constraint.sh" -sh "$TESTS_DIR/sh/test_nvcc_meets_llama_minimum.sh" -sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh" -sh "$TESTS_DIR/sh/test_staged_validation_enabled.sh" -sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh" -sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh" -sh "$TESTS_DIR/sh/test_torch_flavor.sh" -sh "$TESTS_DIR/sh/test_redact_install_output.sh" -sh "$TESTS_DIR/sh/test_install_uv_override_space.sh" -sh "$TESTS_DIR/sh/test_install_rollback_lifecycle.sh" +# Discovered, not listed: a hand-maintained list drifts (this one had fallen +# eight files behind sh/, and the Backend CI copy of it had fallen seven). +# Backend CI discovers the same directory and skips the same file, plus +# test_install_rollback_lifecycle.sh which cross-platform-parity-ci.yml already +# runs on both platforms. tests/studio/test_ci_shell_suite_coverage.py fails if +# either side stops discovering, or skips something undocumented. +# test_install_host_defaults.sh: asserts an install.ps1 layout that has +# drifted (separate followup). +SH_SKIP="test_install_host_defaults.sh" +for _t in "$TESTS_DIR"/sh/test_*.sh; do + case " $SH_SKIP " in + *" $(basename "$_t") "*) echo "skipping $(basename "$_t")"; continue ;; + esac + # bash, not sh: every file under sh/ declares a bash shebang, and three of + # them fail on bashisms under dash, which is /bin/sh on Debian and Ubuntu + # (test_apt_distro_prompt, test_studio_home_node_dir, and + # test_with_llama_cpp_dir_link_behavior). The old hand-written list happened + # to name only dash-clean files, so discovering the directory is what + # exposed this. Backend CI already invokes them with bash. + bash "$_t" +done echo "" echo "=== Python tests ===" diff --git a/tests/studio/install/test_rocm_arch_table_parity.py b/tests/studio/install/test_rocm_arch_table_parity.py new file mode 100644 index 0000000000..40f6858f7e --- /dev/null +++ b/tests/studio/install/test_rocm_arch_table_parity.py @@ -0,0 +1,651 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Drift guards for the AMD gfx tables that are duplicated across the installers. + +The same three tables are hand-copied into up to seven places each: + + gfx -> AMD index family install.sh (_amd_arch_index_family_for_gfx) + install.ps1 ($archFamilyMap) + studio/setup.ps1 ($archFamilyMap) + studio/install_python_stack.py (_GFX_TO_AMD_INDEX_ARCH) + + GPU name -> gfx install.sh (_infer_amd_gfx_arch_from_gpu_name) + install.sh (case "$_gpu_disp_mkt", detection banner + env tip) + studio/setup.sh (case "$_setup_mkt") + install.ps1 ($nameArchTable) + studio/setup.ps1 ($nameArchTable) + studio/install_python_stack.py (_WIN_GPU_NAME_ARCH_TABLE) + tests/_zoo_rocm_spoof.py (_PROFILES, inverted gfx -> name) + + torch>=2.11 pin allowlist install.sh (case "$_torch_index_leaf") + install.ps1 ($_pinGfx211) + studio/setup.ps1 (Test-RocmPinLeaf211) + +Every copy carries a "kept in sync with" comment and nothing enforced it, which is +how the routing family of bugs kept recurring: #7264 / #7280 (Strix left on the +generic rocm7.2 index), #7293 / #7300 (fixed in one installer at a time) and #7277 +(RDNA2 gfx1030-1036 added to install.ps1 / setup.ps1 / install_python_stack.py -- +install.sh had to follow separately). Half-applied edits are invisible until an AMD +user on the missed path gets CPU-only PyTorch. + +These tests parse each copy out of its source file and compare them, so a table +edited in one place fails CI naming the file that was missed. + +Counting the copies by hand is itself unreliable -- the in-code "kept in sync +with" comments claimed four when there were seven -- so TestNoUnregisteredArchTable +below rediscovers them by scanning the repo instead of trusting this list. +""" + +import ast +import fnmatch +import importlib.util +import re +import sys +from pathlib import Path + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +_INSTALL_SH = PACKAGE_ROOT / "install.sh" +_INSTALL_PS1 = PACKAGE_ROOT / "install.ps1" +_SETUP_SH = PACKAGE_ROOT / "studio" / "setup.sh" +_SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1" +_STACK_PY = PACKAGE_ROOT / "studio" / "install_python_stack.py" +_SPOOF_PY = PACKAGE_ROOT / "tests" / "_zoo_rocm_spoof.py" + + +def _load_stack_module(): + spec = importlib.util.spec_from_file_location("studio_install_python_stack_parity", _STACK_PY) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +stack_mod = _load_stack_module() + + +# ── Source extraction helpers ──────────────────────────────────────────────── + + +def _sh_function_body(source: str, name: str) -> str: + """Return a POSIX-shell function body by brace matching (same idea as + _extract_sh_function_body in test_rocm_support.py, kept local so this file + stands alone).""" + needle = f"{name}() {{" + start = source.find(needle) + assert start != -1, f"{name}() not found" + depth = 0 + i = start + len(needle) - 1 + while i < len(source): + if source[i] == "{": + depth += 1 + elif source[i] == "}": + depth -= 1 + if depth == 0: + return source[start : i + 1] + i += 1 + raise AssertionError(f"unterminated {name}()") + + +def _sh_case_block(source: str, subject: str) -> str: + """Return the body of `case in ... esac` (first match).""" + start = source.find(f"case {subject} in") + assert start != -1, f"case {subject} in ... not found" + end = source.find("esac", start) + assert end != -1, f"unterminated case {subject}" + return source[start:end] + + +def _ps_block(source: str, header: str, open_ch: str, close_ch: str) -> str: + """Return the balanced `header ... ` block from a PowerShell file.""" + start = source.find(header) + assert start != -1, f"{header} not found" + i = source.find(open_ch, start) + assert i != -1 + depth = 0 + while i < len(source): + if source[i] == open_ch: + depth += 1 + elif source[i] == close_ch: + depth -= 1 + if depth == 0: + return source[start : i + 1] + i += 1 + raise AssertionError(f"unterminated {header}") + + +def _strip_sh_comment(line: str) -> str: + """Drop a trailing `# ...` comment. Safe here: no table line contains a '#' + inside a pattern.""" + return line.split("#", 1)[0] + + +# ── Table 1: gfx -> AMD index family ───────────────────────────────────────── + + +def _gfx_family_map_sh() -> dict[str, str]: + body = _sh_function_body( + _INSTALL_SH.read_text(encoding = "utf-8"), "_amd_arch_index_family_for_gfx" + ) + out: dict[str, str] = {} + for line in body.splitlines(): + m = re.match(r"\s*(gfx[^)]*)\)\s*echo\s+(\S+)\s*;;", _strip_sh_comment(line)) + if not m: + continue + for arch in m.group(1).split("|"): + out[arch.strip()] = m.group(2).strip() + return out + + +def _gfx_family_map_ps(path: Path) -> dict[str, str]: + block = _ps_block(path.read_text(encoding = "utf-8"), "$archFamilyMap = @{", "{", "}") + out: dict[str, str] = {} + for line in block.splitlines(): + for m in re.finditer(r'"(gfx[0-9a-z]+)"\s*=\s*"([A-Za-z0-9-]+)"', _strip_sh_comment(line)): + out[m.group(1)] = m.group(2) + return out + + +def _gfx_family_maps() -> dict[str, dict[str, str]]: + return { + "studio/install_python_stack.py": dict(stack_mod._GFX_TO_AMD_INDEX_ARCH), + "install.sh": _gfx_family_map_sh(), + "install.ps1": _gfx_family_map_ps(_INSTALL_PS1), + "studio/setup.ps1": _gfx_family_map_ps(_SETUP_PS1), + } + + +class TestGfxIndexFamilyParity: + """All four gfx -> AMD index family maps must agree, entry for entry.""" + + def test_every_copy_is_non_empty(self): + for where, table in _gfx_family_maps().items(): + assert ( + table + ), f"{where}: parsed an empty gfx -> index family map (table moved or renamed?)" + + def test_all_copies_identical(self): + maps = _gfx_family_maps() + reference_name = "studio/install_python_stack.py" + reference = maps[reference_name] + for where, table in maps.items(): + if where == reference_name: + continue + missing = {k: v for k, v in reference.items() if k not in table} + extra = {k: v for k, v in table.items() if k not in reference} + wrong = { + k: (v, reference[k]) + for k, v in table.items() + if k in reference and v != reference[k] + } + assert ( + not missing + ), f"{where} is missing {sorted(missing)} (present in {reference_name})" + assert not extra, f"{where} has {sorted(extra)} that {reference_name} does not" + assert not wrong, f"{where} maps {wrong} (value, expected)" + + def test_rdna2_family_present_everywhere(self): + """#7277 added gfx1030-1036 to three files; install.sh followed later. + Pin the whole RDNA2 range so the next family lands everywhere at once.""" + for where, table in _gfx_family_maps().items(): + for arch in ( + "gfx1030", + "gfx1031", + "gfx1032", + "gfx1033", + "gfx1034", + "gfx1035", + "gfx1036", + ): + assert table.get(arch) == "gfx103X-all", f"{where}: {arch} -> {table.get(arch)!r}" + + +class TestSupportedWheelArchList: + """setup.ps1's $_rocmWheelArches decides whether a detected arch gets ROCm torch + at all. An arch present in the family map but absent here silently installs + CPU-only PyTorch (the 'not in supported arch list' report from r/unsloth).""" + + def test_wheel_arch_list_covers_every_mapped_arch(self): + block = _ps_block( + _SETUP_PS1.read_text(encoding = "utf-8"), "$_rocmWheelArches = @(", "(", ")" + ) + listed = set(re.findall(r'"(gfx[0-9a-z]+)"', block)) + assert listed, "could not parse $_rocmWheelArches" + mapped = set(stack_mod._GFX_TO_AMD_INDEX_ARCH) + assert mapped - listed == set(), ( + f"studio/setup.ps1 $_rocmWheelArches is missing {sorted(mapped - listed)}: " + "those arches map to an AMD index but would still fall back to CPU torch" + ) + + +# ── Table 2: GPU marketing name -> gfx ─────────────────────────────────────── +# +# Each copy is an ordered, first-match-wins table. The shell copies use case +# globs (case-sensitive); the PowerShell copies use -match regexes +# (case-insensitive, and the only place a negative lookahead is available). +# Rather than diff the patterns -- which legitimately differ in syntax -- run +# every copy against the same real GPU names and require the same answer. + + +def _name_table_sh_function(source: str, name: str) -> list[tuple[list[str], str]]: + body = _sh_function_body(source, name) + rows: list[tuple[list[str], str]] = [] + for line in body.splitlines(): + m = re.match(r"\s*(\*.*?)\)\s*echo\s+(gfx[0-9a-z]+)\s*;;", _strip_sh_comment(line)) + if m: + rows.append(([p.strip() for p in m.group(1).split("|")], m.group(2))) + return rows + + +def _name_table_sh_case(source: str, subject: str, var: str) -> list[tuple[list[str], str]]: + """A bare `case ... in` table that assigns to a variable rather than echoing.""" + block = _sh_case_block(source, subject) + rows: list[tuple[list[str], str]] = [] + for line in block.splitlines(): + m = re.match( + rf'\s*(\*.*?)\)\s*{re.escape(var)}="(gfx[0-9a-z]+)"\s*;;', _strip_sh_comment(line) + ) + if m: + rows.append(([p.strip() for p in m.group(1).split("|")], m.group(2))) + return rows + + +def _name_table_ps(path: Path) -> list[tuple[str, str]]: + block = _ps_block(path.read_text(encoding = "utf-8"), "$nameArchTable = @(", "(", ")") + return re.findall(r'@\{\s*P\s*=\s*"([^"]+)"\s*;\s*A\s*=\s*"(gfx[0-9a-z]+)"\s*\}', block) + + +def _match_sh(rows: list[tuple[list[str], str]], gpu_name: str) -> str | None: + """Evaluate a shell `case` table: first arm whose glob matches wins.""" + for patterns, arch in rows: + for pattern in patterns: + # Shell case globs quote literal segments: *"RX 7900"* -> *RX 7900* + if fnmatch.fnmatchcase(gpu_name, pattern.replace('"', "")): + return arch + return None + + +def _match_ps(rows: list[tuple[str, str]], gpu_name: str) -> str | None: + """Evaluate a PowerShell -match table: first arm whose regex matches wins. + -match is case-insensitive; .NET and Python agree on these patterns + (alternation plus one negative lookahead).""" + for pattern, arch in rows: + if re.search(pattern, gpu_name, re.IGNORECASE): + return arch + return None + + +# Real strings as amd-smi / rocm-smi / WMI report them, including the two +# ordering traps: "RX 9070 XT" must beat the bare "9070" arm, and "RX 7700S" +# must beat the "RX 7700" arm. +# +# The expectation is the *AMD pip index leaf*, not the gfx id. The leaf is what +# the tables exist to produce -- it picks the wheel -- and it is what a wrong +# answer actually costs the user. Exact gfx ids are pinned separately in +# _AMD_DOCUMENTED_ARCH, sourced from AMD rather than from these tables. +_GPU_NAME_LEAF_CASES = [ + ("AMD Radeon RX 9070 XT", "gfx120X-all"), + ("AMD Radeon RX 9070", "gfx120X-all"), + ("AMD Radeon RX 9060 XT", "gfx120X-all"), + ("AMD Radeon 8060S Graphics", "gfx1151"), + ("AMD Ryzen AI Max+ 395 w/ Radeon 8060S Graphics", "gfx1151"), + ("AMD Radeon 890M Graphics", "gfx1150"), + ("AMD Radeon 880M Graphics", "gfx1150"), + ("AMD Radeon 860M Graphics", "gfx1152"), + ("AMD Radeon 840M Graphics", "gfx1152"), + ("AMD Ryzen AI 7 350 w/ Radeon 860M", "gfx1152"), + ("AMD Radeon RX 7900 XTX", "gfx110X-all"), + ("AMD Radeon RX 7800 XT", "gfx110X-all"), + ("AMD Radeon PRO W7900", "gfx110X-all"), + ("AMD Radeon RX 7700S", "gfx110X-all"), + ("AMD Radeon RX 7600 XT", "gfx110X-all"), + ("AMD Radeon 780M Graphics", "gfx110X-all"), + ("AMD Radeon RX 6900 XT", "gfx103X-all"), + ("AMD Radeon RX 6700 XT", "gfx103X-all"), + ("AMD Radeon RX 6600 XT", "gfx103X-all"), + ("AMD Radeon RX 6500 XT", "gfx103X-all"), +] + +# Exact gfx ids, transcribed from AMD's ROCm compatibility matrix (the "Radeon +# GPU" list at rocm.docs.amd.com/en/latest/compatibility/compatibility-matrix.html), +# NOT from the installer tables. This is the ground truth the tables are supposed +# to reproduce, so it has to come from outside them. +# +# Three of these were wrong until the commit that added this table: RX 9070 +# (non-XT) said gfx1200, RX 7800 XT / 7700 XT / PRO W7700 said gfx1100, and PRO +# V710 said gfx1102. Nobody was misrouted, because each wrong id happened to +# share an index leaf with the right one, which is exactly why it went unnoticed +# through five copies of the table. The leaf assertions above cannot catch that +# class of error; only an external source can. +# +# The APU rows were added after that: Krackan Point (860M / 840M) said gfx1150 +# but is gfx1152, and unlike the three above that one DID change the wheel, +# since gfx1150 and gfx1152 are separate index leaves on repo.amd.com. +_AMD_DOCUMENTED_ARCH = { + # RDNA 4 -- Navi 48 is gfx1201, Navi 44 is gfx1200. + "AMD Radeon RX 9070 XT": "gfx1201", + "AMD Radeon RX 9070 GRE": "gfx1201", + "AMD Radeon RX 9070": "gfx1201", + "AMD Radeon RX 9060 XT": "gfx1200", + "AMD Radeon RX 9060": "gfx1200", + # RDNA 3 -- Navi 31 / 32 / 33. + "AMD Radeon RX 7900 XTX": "gfx1100", + "AMD Radeon PRO W7900": "gfx1100", + "AMD Radeon PRO W7800": "gfx1100", + "AMD Radeon RX 7800 XT": "gfx1101", + "AMD Radeon RX 7700 XT": "gfx1101", + "AMD Radeon PRO W7700": "gfx1101", + "AMD Radeon PRO V710": "gfx1101", + "AMD Radeon RX 7600 XT": "gfx1102", + "AMD Radeon RX 7700S": "gfx1102", + "AMD Radeon PRO W7600": "gfx1102", + # RDNA 3.5 APUs -- Strix Point is gfx1150, Krackan Point (860M/840M) is + # gfx1152, per AMD's own lemonade GPU table (src/cpp/server/system_info.cpp). + "AMD Radeon 8060S Graphics": "gfx1151", + "AMD Radeon 890M Graphics": "gfx1150", + "AMD Radeon 880M Graphics": "gfx1150", + "AMD Radeon 860M Graphics": "gfx1152", + "AMD Radeon 840M Graphics": "gfx1152", +} + + +def _name_tables() -> dict[str, object]: + install_sh = _INSTALL_SH.read_text(encoding = "utf-8") + return { + "install.sh:_infer_amd_gfx_arch_from_gpu_name": _name_table_sh_function( + install_sh, "_infer_amd_gfx_arch_from_gpu_name" + ), + # install.sh carries the table TWICE. The second copy drives the detection + # banner and, more importantly, the "Tip: set UNSLOTH_ROCM_GFX_ARCH=" + # line, so a wrong id there gets pasted into a user's environment where it + # becomes authoritative. Neither this copy nor the two below were in this + # parity check until the arch-id fix went looking for every place the + # table lives -- six, not four. + "install.sh:_gpu_disp_gfx": _name_table_sh_case( + install_sh, '"$_gpu_disp_mkt"', "_gpu_disp_gfx" + ), + "studio/setup.sh": _name_table_sh_case( + _SETUP_SH.read_text(encoding = "utf-8"), '"$_setup_mkt"', "_setup_gfx" + ), + "install.ps1": _name_table_ps(_INSTALL_PS1), + "studio/setup.ps1": _name_table_ps(_SETUP_PS1), + "studio/install_python_stack.py": list(stack_mod._WIN_GPU_NAME_ARCH_TABLE), + } + + +def _spoof_profiles() -> dict[str, str]: + """gfx -> marketing name out of tests/_zoo_rocm_spoof.py::_PROFILES. + + Parsed with ast rather than imported: that module spoofs torch.cuda and the + AMD identity as an import side effect, which would poison every test sharing + the process.""" + tree = ast.parse(_SPOOF_PY.read_text(encoding = "utf-8")) + for node in tree.body: + target = node.target if isinstance(node, ast.AnnAssign) else None + if target is not None and getattr(target, "id", "") == "_PROFILES": + return {gfx: value[0] for gfx, value in ast.literal_eval(node.value).items()} + raise AssertionError("_PROFILES not found in tests/_zoo_rocm_spoof.py") + + +# The spoof fixture states the mapping backwards (gfx -> the name torch should +# report), so it is the one copy written from the hardware's point of view +# instead of the installer's. That makes it a useful independent witness: it had +# gfx1101 -> "RX 7800 XT" and gfx1201 -> "RX 9070 XT" correct while all six +# installer copies were wrong, and nothing compared the two. +# +# RX 6700 XT is a known, deliberate divergence rather than drift. AMD's +# compatibility matrix documents no consumer RX 6000 card and no gfx1031 at all +# (only "AMD Radeon PRO W6800 (gfx1030)"), the installer arm is commented +# "gfx103X family", and no code consumes the exact id -- gfx1031 appears only as +# a key in the index-family maps, never as a value any name table emits. With no +# external source to correct it against, changing shipped behaviour here would be +# guesswork, so the divergence is pinned instead of silently normalised. +_SPOOF_DIVERGENCES = { + "gfx1031": "installers group Navi 22 into the gfx1030 arm; see comment above", +} + + +def _resolve(where: str, rows, gpu_name: str) -> str | None: + """Shell copies are case globs; the PowerShell and Python copies are both + ordered first-match regex tables evaluated case-insensitively, so _match_ps + models either one. `where` may be ":" for the files that carry + the table more than once.""" + return ( + _match_sh(rows, gpu_name) + if where.split(":")[0].endswith(".sh") + else _match_ps(rows, gpu_name) + ) + + +class TestGpuNameArchParity: + """All four name -> gfx tables must resolve the same GPU the same way.""" + + def test_every_copy_is_non_empty(self): + for where, rows in _name_tables().items(): + assert rows, f"{where}: parsed an empty name -> gfx table (table moved or renamed?)" + + @pytest.mark.parametrize("gpu_name", [name for name, _ in _GPU_NAME_LEAF_CASES]) + def test_all_copies_return_the_same_arch(self, gpu_name): + """The drift guard proper: no expected value, just agreement. This is what + catches a table edited in one installer and not the other three, and it + stays honest even where the shipped gfx id is itself wrong.""" + answers = {where: _resolve(where, rows, gpu_name) for where, rows in _name_tables().items()} + distinct = set(answers.values()) + assert len(distinct) == 1, f"{gpu_name!r} resolves inconsistently: {answers}" + assert distinct != {None}, f"{gpu_name!r} is not matched by any copy of the table" + + @pytest.mark.parametrize("gpu_name,expected_leaf", _GPU_NAME_LEAF_CASES) + def test_every_copy_routes_to_the_right_wheel_index(self, gpu_name, expected_leaf): + """What the tables are for. A wrong leaf is the user-visible failure: + CPU-only torch, or a wheel built for the wrong ISA.""" + families = stack_mod._GFX_TO_AMD_INDEX_ARCH + for where, rows in _name_tables().items(): + arch = _resolve(where, rows, gpu_name) + assert arch is not None, f"{where}: {gpu_name!r} matched nothing" + assert ( + families.get(arch) == expected_leaf + ), f"{where}: {gpu_name!r} -> {arch} -> {families.get(arch)!r}, expected {expected_leaf!r}" + + @pytest.mark.parametrize("gpu_name,expected_arch", sorted(_AMD_DOCUMENTED_ARCH.items())) + def test_every_copy_matches_amds_documented_arch(self, gpu_name, expected_arch): + """The gfx id itself, against AMD's matrix rather than against a sibling + copy of the same table. Agreement between five copies proves nothing if + all five were transcribed from the same mistake.""" + for where, rows in _name_tables().items(): + arch = _resolve(where, rows, gpu_name) + assert ( + arch == expected_arch + ), f"{where}: {gpu_name!r} -> {arch!r}, AMD documents {expected_arch!r}" + + def test_unknown_name_matches_nothing_anywhere(self): + """An unrecognised card must fall through to the CPU path in every copy, + never onto a neighbouring arm.""" + for where, rows in _name_tables().items(): + got = _resolve(where, rows, "NVIDIA GeForce RTX 4090") + assert got is None, f"{where}: RTX 4090 matched {got!r}" + + def test_inferred_arch_always_has_an_index_family(self): + """Every arch a name table can produce must be routable to an AMD wheel + index, else detection succeeds and the install still lands on CPU torch.""" + families = stack_mod._GFX_TO_AMD_INDEX_ARCH + for where, rows in _name_tables().items(): + for arch in {arch for _, arch in rows}: + assert arch in families, f"{where}: {arch} has no entry in _GFX_TO_AMD_INDEX_ARCH" + + def test_every_documented_gpu_resolves_somewhere(self): + """The reverse of the AMD check above. That one asks "do the tables get + the documented cards right"; this asks "is a documented card missing + entirely", which is a silent CPU fallback rather than a wrong id. + + This cannot notice a GPU AMD shipped that nobody transcribed into + _AMD_DOCUMENTED_ARCH -- doing that honestly would mean fetching AMD's + matrix at test time, which makes the suite non-hermetic and offline + runners fail. It does catch a card added to the ground-truth list, or to + one installer, without the tables being completed.""" + for gpu_name in sorted(_AMD_DOCUMENTED_ARCH): + for where, rows in _name_tables().items(): + assert ( + _resolve(where, rows, gpu_name) is not None + ), f"{where}: {gpu_name!r} matches no arm, so this card gets CPU-only torch" + + +class TestSpoofFixtureParity: + """tests/_zoo_rocm_spoof.py is the seventh copy of the name/gfx mapping and + was outside every drift guard. It is the fixture other ROCm tests build their + fake AMD host from, so if it and the installers disagree, those tests exercise + a machine that cannot exist.""" + + def test_spoof_profiles_parse(self): + profiles = _spoof_profiles() + assert profiles, "parsed an empty _PROFILES (renamed or restructured?)" + assert all(gfx.startswith("gfx") for gfx in profiles), profiles + + def test_spoof_names_resolve_back_to_their_own_arch(self): + """Round-trip: feed each spoofed marketing name through the installer + tables and the answer must be the gfx the spoof claims to be emulating.""" + tables = _name_tables() + for gfx, gpu_name in sorted(_spoof_profiles().items()): + if gfx in _SPOOF_DIVERGENCES: + continue + for where, rows in tables.items(): + got = _resolve(where, rows, gpu_name) + assert ( + got == gfx + ), f"{where}: spoof says {gfx} is {gpu_name!r}, installer says {got!r}" + + def test_divergences_are_real_and_still_diverging(self): + """Keeps the exception list from going stale: if the installers are + corrected later, this fails and the entry has to be removed rather than + quietly suppressing a check that now passes.""" + tables = _name_tables() + profiles = _spoof_profiles() + for gfx in _SPOOF_DIVERGENCES: + assert gfx in profiles, f"{gfx} is exempted but no longer in the spoof" + answers = {_resolve(w, r, profiles[gfx]) for w, r in tables.items()} + assert answers != {gfx}, f"{gfx} now agrees everywhere; drop it from _SPOOF_DIVERGENCES" + + +# ── The meta-guard: find copies nobody registered ──────────────────────────── + + +# A table line names a card and gives its arch. Matching both on one line is what +# separates a real table from the many files that merely mention a gfx id (kernel +# dispatch, OOM guards, doc comments). +_MKT_NAME = re.compile(r"(RX\s*\d{4}|PRO\s*[WV]\d{3,4}|\b90[5-8]0\b)", re.IGNORECASE) +_GFX_ID = re.compile(r"gfx1[0-2][0-9a-z]{1,2}") + +# Skip dirs of third-party or generated code; scanning them is slow and any hit +# would not be ours to fix. +_SCAN_SKIP_DIRS = {".git", "node_modules", ".venv", "venv", "build", "dist", "__pycache__"} + +# Every file allowed to carry a name/arch table, as a repo-relative posix path. +# Adding a copy means adding it here AND wiring it into a parity check above; +# that is the point of the guard. +_REGISTERED_TABLE_FILES = { + "install.sh", + "install.ps1", + "studio/setup.sh", + "studio/setup.ps1", + "studio/install_python_stack.py", + "tests/_zoo_rocm_spoof.py", +} + +# Three or more such lines means a table. One or two means prose: the two known +# single-line hits are comments ("Verified on gfx1151 (Radeon 8060S)" in +# scripts/install_rocm_wsl_strixhalo.sh, and a parenthetical in +# studio/install_llama_prebuilt.py). Real tables score 9 to 17, so the gap is +# wide and the threshold is not load-bearing. +_TABLE_LINE_THRESHOLD = 3 + + +def _files_carrying_a_name_arch_table() -> dict[str, int]: + found: dict[str, int] = {} + for path in PACKAGE_ROOT.rglob("*"): + if path.suffix not in {".sh", ".ps1", ".py"} or not path.is_file(): + continue + rel = path.relative_to(PACKAGE_ROOT).as_posix() + if any(part in _SCAN_SKIP_DIRS for part in path.relative_to(PACKAGE_ROOT).parts): + continue + # Tests that *assert* on the tables quote card names next to gfx ids by + # nature. Fixtures like _zoo_rocm_spoof.py do not start with test_ and so + # stay in scope, which is how the seventh copy surfaced. + if path.name.startswith("test_"): + continue + try: + text = path.read_text(encoding = "utf-8", errors = "ignore") + except OSError: + continue + hits = sum( + 1 for line in text.splitlines() if _MKT_NAME.search(line) and _GFX_ID.search(line) + ) + if hits >= _TABLE_LINE_THRESHOLD: + found[rel] = hits + return found + + +class TestNoUnregisteredArchTable: + """The failure this whole file exists for is a copy of the table that nobody + knew about. Enumerating the copies by hand is the same manual step that let + them drift, so this rediscovers them from the source tree.""" + + def test_scan_still_finds_the_known_copies(self): + """Guards the guard: if the heuristic stops matching (patterns reformatted + onto multiple lines, say), it would silently find nothing and pass.""" + found = _files_carrying_a_name_arch_table() + missing = _REGISTERED_TABLE_FILES - set(found) + assert not missing, f"scan no longer detects known tables in {sorted(missing)}" + + def test_no_unregistered_copies(self): + found = _files_carrying_a_name_arch_table() + extra = {rel: n for rel, n in found.items() if rel not in _REGISTERED_TABLE_FILES} + assert not extra, ( + f"unregistered GPU-name/arch table(s): {extra}. Wire each into " + f"_name_tables() (or the spoof check) and add it to " + f"_REGISTERED_TABLE_FILES, so drift there fails CI too." + ) + + +# ── Table 3: the torch>=2.11 pin allowlist ─────────────────────────────────── + + +class TestTorch211PinAllowlistParity: + """gfx120X-all / gfx1151 / gfx1150 / gfx1152 (and rocm7.2) ship the null + _grouped_mm kernel below torch 2.11, so all three installers must raise the + same floor. A leaf missing from one copy reintroduces the crash there.""" + + _EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150", "gfx1152"} + + def test_install_sh_pins_the_same_leaves(self): + source = _INSTALL_SH.read_text(encoding = "utf-8") + idx = source.find('case "$_torch_index_leaf" in') + assert idx != -1 + arm = re.search(r"\n\s*(rocm7\.2\|[^)]*)\)", source[idx:]) + assert arm, "torch 2.11 pin arm not found in install.sh" + leaves = {leaf.strip() for leaf in arm.group(1).split("|")} + assert ( + self._EXPECTED <= leaves + ), f"install.sh pin arm missing {sorted(self._EXPECTED - leaves)}" + assert "rocm7.2" in leaves + + def test_install_ps1_pins_the_same_leaves(self): + source = _INSTALL_PS1.read_text(encoding = "utf-8") + m = re.search(r"\$_pinGfx211\s*=\s*@\(([^)]*)\)", source) + assert m, "$_pinGfx211 not found in install.ps1" + leaves = set(re.findall(r"'([^']+)'", m.group(1))) + assert leaves == self._EXPECTED, f"install.ps1 pins {sorted(leaves)}" + + def test_setup_ps1_pins_the_same_leaves(self): + source = _SETUP_PS1.read_text(encoding = "utf-8") + m = re.search(r"return\s+@\(([^)]*)\)\s*-contains\s*\$Leaf", source) + assert m, "the 2.11 pin allowlist helper was not found in studio/setup.ps1" + leaves = set(re.findall(r"'([^']+)'", m.group(1))) + assert leaves == self._EXPECTED, f"studio/setup.ps1 pins {sorted(leaves)}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/studio/install/test_rocm_native_linux_lib_dirs.py b/tests/studio/install/test_rocm_native_linux_lib_dirs.py new file mode 100644 index 0000000000..9b95af88b0 --- /dev/null +++ b/tests/studio/install/test_rocm_native_linux_lib_dirs.py @@ -0,0 +1,438 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Coverage for the native-Linux system-ROCm library prepend (PR #7233). + +#7233 fixed the segfault-on-launch class of AMD reports (#7208, #7310, #6276 and +the native-Linux half of #7307): a prebuilt llama.cpp ships its own libggml-hip / +HIP runtime, and on a bare-metal ROCm box that bundled runtime can disagree with +the host amdkfd driver, so the server dies the moment a model is loaded. The fix +prepends the *system* ROCm lib dirs ahead of the bundle on LD_LIBRARY_PATH. + +It landed as two hand-copied helpers, one in the installer (validation-time) and +one in the serve-time launcher: + + studio/install_llama_prebuilt.py _bundled_hip_present / _native_linux_system_rocm_lib_dirs + studio/backend/core/inference/llama_cpp.py same two, "mirrors" comment only + +and shipped with no tests at all: the WSL sibling helper added earlier has +TestWslSystemRocmLibDirs / TestBinaryEnvWslOrdering / TestLlamaCppRuntimeWslOrdering, +the native-Linux one has nothing. Every gate here is a false-positive risk that +would silently reorder LD_LIBRARY_PATH for users the fix was never meant to touch +(WSL, NVIDIA hosts, macOS, containers without /dev/kfd), so each gate gets a test, +and both copies are run against the same fake host and required to agree. + +llama_cpp.py cannot be imported from the test suite (module-level structlog / +backend imports), so its two helpers are lifted out with ast and exec'd standalone. +""" + +import ast +import importlib.util +import os +import sys +from pathlib import Path +from unittest.mock import patch + +import pytest + + +PACKAGE_ROOT = Path(__file__).resolve().parents[3] + +_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py" +_LLAMA_CPP_PATH = PACKAGE_ROOT / "studio" / "backend" / "core" / "inference" / "llama_cpp.py" + +_HELPERS = ("_bundled_hip_present", "_native_linux_system_rocm_lib_dirs") + + +def _load_prebuilt_module(): + spec = importlib.util.spec_from_file_location( + "studio_install_llama_prebuilt_native", _PREBUILT_PATH + ) + assert spec is not None and spec.loader is not None + mod = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = mod + spec.loader.exec_module(mod) + return mod + + +def _extract_functions(path: Path, names) -> dict: + """exec just the named top-level functions out of a module that is too + heavy to import.""" + tree = ast.parse(path.read_text(encoding = "utf-8")) + wanted = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names] + found = {n.name for n in wanted} + assert found == set(names), f"{path.name}: missing {sorted(set(names) - found)}" + module = ast.Module(body = wanted, type_ignores = []) + ns: dict = {"os": os, "sys": sys, "Path": Path} + exec(compile(module, str(path), "exec"), ns) + return ns + + +prebuilt_mod = _load_prebuilt_module() +llama_ns = _extract_functions(_LLAMA_CPP_PATH, _HELPERS) + + +def _impls(): + """The two copies of the helper, by the file they live in.""" + return { + "studio/install_llama_prebuilt.py": prebuilt_mod._native_linux_system_rocm_lib_dirs, + "studio/backend/core/inference/llama_cpp.py": llama_ns[ + "_native_linux_system_rocm_lib_dirs" + ], + } + + +def _norm(paths): + """os.path.join emits '\\' on the Windows test host; compare POSIX-style.""" + return [str(p).replace("\\", "/") for p in paths] + + +def _fake_exists(present): + """os.path.exists stub over a set of POSIX paths.""" + present = {p.replace("\\", "/") for p in present} + + def _exists(p): + return str(p).replace("\\", "/") in present + + return _exists + + +@pytest.fixture +def bundle_dir(tmp_path): + """A prebuilt directory that does contain a bundled HIP runtime.""" + d = tmp_path / "bundle" + d.mkdir() + (d / "libggml-hip.so").write_text("") + return d + + +@pytest.fixture(autouse = True) +def _clean_rocm_env(monkeypatch): + for var in ("UNSLOTH_LLAMA_NO_SYSTEM_ROCM", "HIP_PATH", "HIP_PATH_57", "ROCM_PATH"): + monkeypatch.delenv(var, raising = False) + + +def _call( + impl, + bundle, + present, + platform = "linux", +): + """Run one copy of the helper against a fake host. + + sys.platform is patched inside the call rather than in a fixture: pytest's own + tmp_path factory branches on it, so a session-wide patch breaks the fixture on + a Windows test host.""" + with patch.object(sys, "platform", platform): + with patch("os.path.exists", _fake_exists(present)): + return _norm(impl(str(bundle))) + + +class TestBundledHipPresent: + """The prepend only makes sense when the prebuilt actually bundles HIP; a + CPU or CUDA build must be left alone.""" + + @pytest.mark.parametrize( + "where", ["studio/install_llama_prebuilt.py", "studio/backend/core/inference/llama_cpp.py"] + ) + def test_detects_versioned_and_plain_sonames(self, tmp_path, where): + impl = ( + prebuilt_mod._bundled_hip_present + if where == "studio/install_llama_prebuilt.py" + else llama_ns["_bundled_hip_present"] + ) + plain = tmp_path / "plain" + plain.mkdir() + (plain / "libggml-hip.so").write_text("") + versioned = tmp_path / "versioned" + versioned.mkdir() + (versioned / "libggml-hip.so.0.0.1").write_text("") + cpu_only = tmp_path / "cpu" + cpu_only.mkdir() + (cpu_only / "libggml-cpu.so").write_text("") + assert impl(str(plain)) is True, f"{where}: plain soname not detected" + assert impl(str(versioned)) is True, f"{where}: versioned soname not detected" + assert impl(str(cpu_only)) is False, f"{where}: CPU-only build treated as HIP" + assert impl("") is False, f"{where}: empty binary_dir must be falsy" + assert ( + impl(str(tmp_path / "does-not-exist")) is False + ), f"{where}: missing dir must be falsy" + + +class TestNativeLinuxGates: + """Each gate, on both copies. A gate that stops working silently reorders + LD_LIBRARY_PATH for a platform the fix was never aimed at.""" + + _ROCM_LIB = "/opt/rocm/lib" + _HOST = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so"} + + _run = staticmethod(_call) + + def test_returns_system_rocm_lib_on_a_native_rocm_host(self, bundle_dir): + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, self._HOST) == [self._ROCM_LIB], where + + def test_accepts_versioned_hsa_runtime_soname(self, bundle_dir): + present = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so.1"} + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == [self._ROCM_LIB], where + + @pytest.mark.parametrize("platform", ["win32", "darwin"]) + def test_no_op_off_linux(self, bundle_dir, platform): + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, self._HOST, platform = platform) == [], where + + def test_no_op_on_wsl(self, bundle_dir): + """WSL has its own ordering path (plus HSA_ENABLE_DXG_DETECTION); /dev/dxg + must hand off to it, not double-prepend here.""" + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, self._HOST | {"/dev/dxg"}) == [], where + + def test_no_op_without_amdkfd(self, bundle_dir): + """No /dev/kfd: NVIDIA host, CPU host, or a container without the AMD + device node. Prepending system ROCm there would be pure breakage.""" + present = {"/opt/rocm/lib/libhsa-runtime64.so"} + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == [], where + + def test_no_op_when_prebuilt_bundles_no_hip(self, tmp_path): + cpu_bundle = tmp_path / "cpu-bundle" + cpu_bundle.mkdir() + for where, impl in _impls().items(): + assert self._run(impl, cpu_bundle, self._HOST) == [], where + + def test_no_op_when_system_rocm_has_no_hsa_runtime(self, bundle_dir): + """ROCm dir exists but is not a usable runtime install.""" + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, {"/dev/kfd"}) == [], where + + def test_opt_out_env_wins_over_everything(self, bundle_dir, monkeypatch): + monkeypatch.setenv("UNSLOTH_LLAMA_NO_SYSTEM_ROCM", "1") + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, self._HOST) == [], where + + def test_opt_out_env_only_honours_exactly_one(self, bundle_dir, monkeypatch): + """Documented switch is =1; "0"/"" must not disable the fix.""" + for value in ("0", "", "false"): + monkeypatch.setenv("UNSLOTH_LLAMA_NO_SYSTEM_ROCM", value) + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, self._HOST) == [ + self._ROCM_LIB + ], f"{where} ({value!r})" + + +class TestNativeLinuxRootResolution: + """Which ROCm roots are searched, in what order.""" + + _run = staticmethod(_call) + + def test_env_roots_take_precedence_over_opt_rocm(self, bundle_dir, monkeypatch): + """A user with a side-by-side ROCm (HIP_PATH) must get theirs first: the + one matching their driver, not whatever /opt/rocm happens to be.""" + monkeypatch.setenv("HIP_PATH", "/usr/local/rocm7") + present = { + "/dev/kfd", + "/usr/local/rocm7/lib/libhsa-runtime64.so", + "/opt/rocm/lib/libhsa-runtime64.so", + } + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == [ + "/usr/local/rocm7/lib", + "/opt/rocm/lib", + ], where + + def test_all_three_env_roots_are_consulted_in_order(self, bundle_dir, monkeypatch): + monkeypatch.setenv("HIP_PATH", "/a") + monkeypatch.setenv("HIP_PATH_57", "/b") + monkeypatch.setenv("ROCM_PATH", "/c") + present = { + "/dev/kfd", + "/a/lib/libhsa-runtime64.so", + "/b/lib/libhsa-runtime64.so", + "/c/lib/libhsa-runtime64.so", + } + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == ["/a/lib", "/b/lib", "/c/lib"], where + + def test_lib64_layout_is_found(self, bundle_dir): + """RHEL / SUSE ROCm packages install to lib64.""" + present = {"/dev/kfd", "/opt/rocm/lib64/libhsa-runtime64.so"} + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == ["/opt/rocm/lib64"], where + + def test_lib_precedes_lib64_when_both_exist(self, bundle_dir): + present = { + "/dev/kfd", + "/opt/rocm/lib/libhsa-runtime64.so", + "/opt/rocm/lib64/libhsa-runtime64.so", + } + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == [ + "/opt/rocm/lib", + "/opt/rocm/lib64", + ], where + + def test_duplicate_roots_are_deduped(self, bundle_dir, monkeypatch): + """ROCM_PATH=/opt/rocm is the common setup; it must not emit the dir twice.""" + monkeypatch.setenv("ROCM_PATH", "/opt/rocm") + present = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so"} + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == ["/opt/rocm/lib"], where + + def test_empty_env_var_is_ignored(self, bundle_dir, monkeypatch): + monkeypatch.setenv("HIP_PATH", "") + present = {"/dev/kfd", "/opt/rocm/lib/libhsa-runtime64.so"} + for where, impl in _impls().items(): + assert self._run(impl, bundle_dir, present) == ["/opt/rocm/lib"], where + + +class TestHelperParity: + """The two copies carry a "mirrors ..." comment and nothing enforced it.""" + + @pytest.mark.parametrize("name", _HELPERS) + def test_bodies_are_identical(self, name): + a = _function_ast(_PREBUILT_PATH, name) + b = _function_ast(_LLAMA_CPP_PATH, name) + assert ast.dump(a) == ast.dump(b), ( + f"{name} has drifted between install_llama_prebuilt.py and llama_cpp.py; " + "the install-time and serve-time launchers must resolve the same lib dirs" + ) + + +def _function_ast(path: Path, name: str) -> ast.FunctionDef: + """The function's executable body, with docstring and type annotations + stripped: llama_cpp.py quotes its annotations ('list[str]') for the + older-typing lint and documents itself as mirroring the installer. Neither is + drift; the code is.""" + tree = ast.parse(path.read_text(encoding = "utf-8")) + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == name: + for child in ast.walk(node): + if isinstance(child, (ast.AnnAssign, ast.arg)): + child.annotation = None + elif isinstance(child, ast.FunctionDef): + child.returns = None + if ( + node.body + and isinstance(node.body[0], ast.Expr) + and isinstance(node.body[0].value, ast.Constant) + and isinstance(node.body[0].value.value, str) + ): + node.body = node.body[1:] + return node + raise AssertionError(f"{name} not found in {path.name}") + + +class TestBinaryEnvNativeOrdering: + """install-time validation launches the binary through binary_env.""" + + @staticmethod + def _linux_host(): + return prebuilt_mod.HostInfo( + system = "Linux", + machine = "x86_64", + is_windows = False, + is_linux = True, + is_macos = False, + is_x86_64 = True, + is_arm64 = False, + nvidia_smi = None, + driver_cuda_version = None, + compute_caps = [], + visible_cuda_devices = None, + has_physical_nvidia = False, + has_usable_nvidia = False, + has_rocm = True, + ) + + def test_system_rocm_precedes_bundle_dir(self, tmp_path): + binary = tmp_path / "bundle" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("") + sys_rocm = tmp_path / "sysrocm" # dedupe_existing_dirs drops missing dirs + sys_rocm.mkdir() + with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []): + with patch.object( + prebuilt_mod, "_native_linux_system_rocm_lib_dirs", return_value = [str(sys_rocm)] + ): + with patch.dict(os.environ, {}, clear = True): + env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) + ld = [str(Path(p).resolve()) for p in env["LD_LIBRARY_PATH"].split(os.pathsep)] + assert ld.index(str(sys_rocm.resolve())) < ld.index(str(binary.parent.resolve())) + + def test_native_path_does_not_enable_dxg_detection(self, tmp_path): + """HSA_ENABLE_DXG_DETECTION belongs to the WSL branch only; setting it on + bare metal changes HSA agent enumeration for every native AMD user.""" + binary = tmp_path / "bundle" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("") + sys_rocm = tmp_path / "sysrocm" + sys_rocm.mkdir() + with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []): + with patch.object( + prebuilt_mod, "_native_linux_system_rocm_lib_dirs", return_value = [str(sys_rocm)] + ): + with patch.dict(os.environ, {}, clear = True): + env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) + assert "HSA_ENABLE_DXG_DETECTION" not in env + + def test_helper_is_asked_about_the_binary_dir_not_the_install_dir(self, tmp_path): + """_bundled_hip_present globs the directory it is handed; passing + install_dir would look for libggml-hip.so in the wrong place and no-op.""" + binary = tmp_path / "bundle" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("") + seen = [] + + def _spy(binary_dir = ""): + seen.append(binary_dir) + return [] + + with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []): + with patch.object(prebuilt_mod, "_native_linux_system_rocm_lib_dirs", _spy): + with patch.dict(os.environ, {}, clear = True): + prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) + assert seen == [str(binary.parent)] + + def test_no_prepend_leaves_bundle_dir_first(self, tmp_path): + binary = tmp_path / "bundle" / "llama-server" + binary.parent.mkdir(parents = True) + binary.write_text("") + with patch.object(prebuilt_mod, "_wsl_system_rocm_lib_dirs", return_value = []): + with patch.object(prebuilt_mod, "_native_linux_system_rocm_lib_dirs", return_value = []): + with patch.dict(os.environ, {}, clear = True): + env = prebuilt_mod.binary_env(binary, tmp_path, self._linux_host()) + assert env["LD_LIBRARY_PATH"].split(os.pathsep)[0] == str(binary.parent) + + +class TestLlamaCppRuntimeNativeOrdering: + """The serve-time launcher builds LD_LIBRARY_PATH inline inside a large + function, so this half stays a source check (as the WSL sibling does).""" + + def test_prepends_before_binary_dir(self): + source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") + idx_helper = source.find("lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir))") + idx_binary = source.find("lib_dirs.append(binary_dir)") + assert ( + idx_helper != -1 + ), "serve-time launcher must call the native-Linux helper with binary_dir" + assert idx_binary != -1 + assert ( + idx_helper < idx_binary + ), "system ROCm must be searched before the bundled HIP runtime" + + def test_dxg_detection_stays_on_the_wsl_branch(self): + """HSA_ENABLE_DXG_DETECTION must be set from the WSL helper's result only.""" + source = _LLAMA_CPP_PATH.read_text(encoding = "utf-8") + idx_wsl = source.find("lib_dirs.extend(_wsl_system_rocm_lib_dirs())") + idx_dxg = source.find('env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")', idx_wsl) + idx_native = source.find("lib_dirs.extend(_native_linux_system_rocm_lib_dirs(binary_dir))") + assert idx_wsl != -1 and idx_dxg != -1 and idx_native != -1 + assert idx_wsl < idx_dxg < idx_native, ( + "HSA_ENABLE_DXG_DETECTION must be decided from the WSL dirs alone, before the " + "native-Linux dirs are appended to lib_dirs" + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/studio/install/test_rocm_support.py b/tests/studio/install/test_rocm_support.py index b919206120..4f65857a3e 100644 --- a/tests/studio/install/test_rocm_support.py +++ b/tests/studio/install/test_rocm_support.py @@ -2695,9 +2695,11 @@ class TestGfxArchNameFallback: ("AMD Radeon(TM) 890M", "gfx1150"), ("AMD Ryzen AI 9 HX 370 w/ Radeon 890M", "gfx1150"), ("AMD Radeon RX 9070 XT", "gfx1201"), - ("AMD Radeon RX 9070", "gfx1200"), - ("AMD Radeon RX 7700S", "gfx1102"), # (?!S) lookahead must not hit gfx1100 - ("AMD Radeon RX 7700 XT", "gfx1100"), + ("AMD Radeon RX 9070", "gfx1201"), # Navi 48 like the XT, not Navi 44 + ("AMD Radeon RX 9060 XT", "gfx1200"), # Navi 44 + ("AMD Radeon RX 7700S", "gfx1102"), # (?!S) lookahead must not hit gfx1101 + ("AMD Radeon RX 7700 XT", "gfx1101"), # Navi 32 + ("AMD Radeon RX 7900 XTX", "gfx1100"), # Navi 31 ("AMD Radeon(TM) 780M", "gfx1103"), ("NVIDIA GeForce RTX 4090", None), ("Microsoft Basic Display Adapter", None), @@ -3845,7 +3847,7 @@ class TestStrixRocm71Override: assert m._infer_linux_amd_gfx_arch() == "gfx1151" def test_install_sh_cpuinfo_inference_requires_pci_evidence(self): - """install.sh mirror of the VM/container guard: both cpuinfo greps must be + """install.sh mirror of the VM/container guard: every cpuinfo grep must be gated on _gpu_evidence (AMD PCI display device via _amd_gpu_present_via_pci, or the WSL librocdxg gate), and the gate must sit before the first grep.""" source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") @@ -3855,9 +3857,9 @@ class TestStrixRocm71Override: infer = body.find("grep -qiE 'Ryzen AI Max") assert pci >= 0 and infer >= 0 assert pci < infer, "the PCI evidence check must run before the cpuinfo inference" - assert ( - body.count('[ -n "$_gpu_evidence" ] && grep -qiE') == 2 - ), "both cpuinfo greps (gfx1151 and gfx1150) must be gated on _gpu_evidence" + assert body.count("grep -qiE") == body.count( + '[ -n "$_gpu_evidence" ] && grep -qiE' + ), "every cpuinfo grep (gfx1151/gfx1150/gfx1152) must be gated on _gpu_evidence" def test_lspci_scan_covers_all_display_controllers(self): """The lspci fallback must scan every display-class line, not just the @@ -4211,11 +4213,14 @@ class TestStrixRocm71Override: source = _INSTALL_SH_PATH.read_text(encoding = "utf-8") # The 2.11 constraint block must switch on $_torch_index_leaf, not the full # $TORCH_INDEX_URL (a */gfx* match false-positives on a mirror base path). Only the - # _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150) are pushed to 2.11; + # _grouped_mm-bug gfx families (gfx120X-all / gfx1151 / gfx1150 / gfx1152) go to 2.11; # a bare gfx* would also floor gfx110X-all/gfx90a/gfx908, left bare on purpose. - assert 'case "$_torch_index_leaf" in\n rocm7.2|gfx120x-all|gfx1151|gfx1150)' in source, ( + assert ( + 'case "$_torch_index_leaf" in\n rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152)' + in source + ), ( "the torch>=2.11 constraint must match the specific gfx leaves that need " - "it (rocm7.2|gfx120x-all|gfx1151|gfx1150), not a bare gfx* or the whole URL" + "it (rocm7.2|gfx120x-all|gfx1151|gfx1150|gfx1152), not a bare gfx* or the URL" ) def test_amd_rocm_mirror_env_var_respected(self): diff --git a/tests/studio/test_ci_shell_suite_coverage.py b/tests/studio/test_ci_shell_suite_coverage.py new file mode 100644 index 0000000000..50ed22f9de --- /dev/null +++ b/tests/studio/test_ci_shell_suite_coverage.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Guards that the installer test suites actually run on a PR. + +Two ways coverage went missing without anyone noticing: + +1. Backend CI ran a hardcoded list of tests/sh/*.sh files. New tests were added + to the directory and never to the list, so by the time this was written the + list was seven files behind -- including test_strixhalo_wsl_reroute.sh, the + only shell coverage of the ROCm WSL reroute, which had never run on a PR. + tests/run_all.sh, the local entrypoint, had drifted the other way. + +2. Backend CI's path filter did not include install.sh / install.ps1, while a + large share of the suites it runs (tests/sh/*, tests/studio/install/*) assert + against exactly those two files. An install-only change -- the shape most + AMD/ROCm routing fixes take, e.g. #7277 / #7293 / #7300 -- skipped the + workflow that tests it. + +Both are now discovery-based. These tests fail if either reverts to a list, if a +shell test lands somewhere the discovery cannot see it, or if a skip is added +without a reason next to it. +""" + +import re +from pathlib import Path + +import pytest +import yaml + + +REPO_ROOT = Path(__file__).resolve().parents[2] + +_WORKFLOWS = REPO_ROOT / ".github" / "workflows" +_BACKEND_CI = _WORKFLOWS / "studio-backend-ci.yml" +_PARITY_CI = _WORKFLOWS / "cross-platform-parity-ci.yml" +_RUN_ALL = REPO_ROOT / "tests" / "run_all.sh" +_SH_DIR = REPO_ROOT / "tests" / "sh" + +# Files deliberately not run by the auto-discovered Backend CI step. Each needs +# a reason here AND in the workflow; anything else in tests/sh must run. +_EXPECTED_CI_SKIPS = { + "test_install_host_defaults.sh": "asserts an install.ps1 layout that has drifted", + "test_install_rollback_lifecycle.sh": "runs on both platforms in cross-platform-parity-ci.yml", +} + + +def _backend_ci() -> dict: + return yaml.safe_load(_BACKEND_CI.read_text(encoding = "utf-8")) + + +def _shell_step_script() -> str: + """The `run:` body of the shell-installer step, located by name through the + parsed YAML rather than by slicing the raw file.""" + for job in _backend_ci()["jobs"].values(): + for step in job.get("steps", []): + if step.get("name") == "Shell installer tests": + return step["run"] + raise AssertionError("Backend CI has no 'Shell installer tests' step") + + +def _shell_test_files(): + files = sorted(p.name for p in _SH_DIR.glob("test_*.sh")) + assert files, "tests/sh has no test_*.sh files -- did the directory move?" + return files + + +def _skip_list(source: str) -> set[str]: + """The skip= / SH_SKIP= line from a discovery loop.""" + m = re.search(r"^\s*(?:skip|SH_SKIP)=\"([^\"]*)\"", source, re.MULTILINE) + assert m, "no skip list found; the discovery loop must declare one (even if empty)" + return {name for name in m.group(1).split() if name} + + +class TestBackendCiRunsEveryShellTest: + def test_step_discovers_the_directory_instead_of_listing_files(self): + """Matched against the parsed step script, and on the glob rather than a + verbatim line, so reformatting the loop does not turn CI red -- only + going back to a hardcoded list does.""" + script = _shell_step_script() + assert re.search(r"for\s+\w+\s+in\s+tests/sh/test_\*\.sh", script), ( + "Backend CI must glob tests/sh; a hardcoded list is how the ROCm WSL " + f"suite went unrun for months. Step script was:\n{script}" + ) + listed = re.findall(r"tests/sh/test_[a-z0-9_]+\.sh", script) + assert not listed, f"Backend CI still names individual shell tests: {sorted(set(listed))}" + + def test_step_fails_loudly_if_discovery_finds_nothing(self): + """A moved directory must break the build, not pass vacuously.""" + assert "no shell tests discovered under tests/sh" in _shell_step_script() + + def test_every_shell_test_runs_or_is_a_known_skip(self): + skips = _skip_list(_shell_step_script()) + unexpected = skips - set(_EXPECTED_CI_SKIPS) + assert not unexpected, ( + f"Backend CI skips {sorted(unexpected)} without a reason recorded in " + "_EXPECTED_CI_SKIPS; add one or stop skipping it" + ) + # Everything else in the directory is covered by the glob. + for name in _shell_test_files(): + assert name not in skips or name in _EXPECTED_CI_SKIPS, name + + def test_skip_entries_are_not_stale(self): + """A skip for a deleted file quietly widens next time a name is reused.""" + existing = set(_shell_test_files()) + for name in _skip_list(_shell_step_script()): + assert name in existing, f"{name} is skipped but no longer exists in tests/sh" + + def test_each_skip_is_documented_in_the_workflow(self): + source = _BACKEND_CI.read_text(encoding = "utf-8") + for name in _EXPECTED_CI_SKIPS: + assert ( + source.count(name) >= 2 + ), f"{name} is skipped in Backend CI without a comment explaining why" + + def test_rollback_lifecycle_really_does_run_elsewhere(self): + """The one skip justified by 'another workflow covers it' must be true.""" + assert "tests/sh/test_install_rollback_lifecycle.sh" in _PARITY_CI.read_text( + encoding = "utf-8" + ) + + def test_rocm_shell_suite_is_in_scope(self): + """The suite whose absence prompted this file: it must exist and be + picked up (i.e. not skipped).""" + assert "test_strixhalo_wsl_reroute.sh" in _shell_test_files() + assert "test_strixhalo_wsl_reroute.sh" not in _skip_list(_shell_step_script()) + + +class TestRunAllMatchesCi: + """tests/run_all.sh is what a contributor runs before pushing. If it and CI + disagree, one of them is lying about the state of the tree.""" + + def test_run_all_discovers_the_directory(self): + source = _RUN_ALL.read_text(encoding = "utf-8") + assert 'for _t in "$TESTS_DIR"/sh/test_*.sh; do' in source + + def test_run_all_invokes_the_tests_with_bash(self): + """Both runners must use the interpreter the tests declare. Every file + under tests/sh/ has a bash shebang, and on Debian/Ubuntu /bin/sh is + dash, under which three of them fail on bashisms. Running them with sh + would fail the suite locally for reasons CI never reproduces.""" + source = _RUN_ALL.read_text(encoding = "utf-8") + assert 'bash "$_t"' in source, "tests/run_all.sh must run tests/sh/ with bash" + assert 'sh "$_t"' not in source.replace( + 'bash "$_t"', "" + ), "tests/run_all.sh still invokes a discovered test with sh" + assert 'bash "$s"' in _shell_step_script(), "Backend CI must run tests/sh/ with bash" + + def test_run_all_skips_are_a_subset_of_ci_skips(self): + local = _skip_list(_RUN_ALL.read_text(encoding = "utf-8")) + unexpected = local - set(_EXPECTED_CI_SKIPS) + assert not unexpected, ( + f"tests/run_all.sh skips {sorted(unexpected)} that CI still runs: a " + "contributor would see green locally and red on the PR" + ) + + +class TestBackendCiPathFilters: + """The workflow has to fire on the files its tests assert against.""" + + def _paths(self) -> set[str]: + """Read the real trigger through the YAML parser. `on:` is a YAML 1.1 + boolean, so pyyaml keys it as True.""" + wf = _backend_ci() + triggers = wf.get("on", wf.get(True)) + assert triggers, "Backend CI has no trigger block" + paths = triggers["pull_request"]["paths"] + assert paths, "Backend CI pull_request trigger has no paths filter" + return set(paths) + + @pytest.mark.parametrize( + "path,why", + [ + ("install.sh", "tests/sh/* and tests/studio/install/* assert against it"), + ("install.ps1", "the Windows/ROCm arch tables and pin allowlist live here"), + ("studio/**", "covers studio/setup.sh, studio/setup.ps1, install_python_stack.py"), + ("tests/**", "test-only changes must run the tests they touch"), + ], + ) + def test_trigger_covers(self, path, why): + assert path in self._paths(), f"Backend CI does not run when {path} changes ({why})" + + def test_installer_change_would_trigger_the_workflow(self): + """End to end: the exact filenames the ROCm fixes edit.""" + paths = self._paths() + for changed in ("install.sh", "install.ps1"): + assert changed in paths + for changed in ("studio/setup.ps1", "studio/setup.sh", "studio/install_python_stack.py"): + assert any( + changed.startswith(pattern.rstrip("*").rstrip("/")) + for pattern in paths + if pattern.endswith("/**") + ), f"nothing in the path filter matches {changed}" + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) From 97475be347cdbda00e2a563ef8da9694090afaf4 Mon Sep 17 00:00:00 2001 From: alkinun Date: Sun, 26 Jul 2026 04:53:00 +0300 Subject: [PATCH 08/25] fix(studio): support hostname-based enterprise proxies (#7416) * fix(studio): support hostname-based enterprise proxies * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix(studio): strip userinfo from proxy fetch targets --------- 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/core/inference/tools.py | 22 ++++--- studio/backend/run.py | 10 ++++ .../backend/tests/test_secure_tunnel_gate.py | 8 +++ .../tests/test_web_fetch_extraction.py | 57 +++++++++++++++++++ tests/studio/test_cli_studio_defaults.py | 7 +++ unsloth_cli/commands/studio.py | 30 ++++++++++ 6 files changed, 126 insertions(+), 8 deletions(-) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index 0ef6dd46cf..b31db3faf6 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -49,6 +49,7 @@ from loggers import get_logger logger = get_logger(__name__) _EXEC_TIMEOUT = 300 # 5 minutes +_DISABLE_DNS_PINNING_ENV = "UNSLOTH_STUDIO_DISABLE_DNS_PINNING" # Splits the UI source-map from the result; loops strip it (like __IMAGES__). RAG_SOURCES_SENTINEL = "\n__RAG_SOURCES__:" @@ -4194,13 +4195,18 @@ def _fetch_url_raw( budget_error = _fetch_budget_exceeded(deadline, cancel_event) if budget_error is not None: return budget_error, "", "" - # Pin to the validated IP (prevents DNS rebinding): rewrite URL to - # the IP, set the Host header. cp = urlparse(current_url) - # Bracket IPv6 addresses so the netloc is valid in a URL. - ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip - ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str - pinned_url = urlunparse(cp._replace(netloc = ip_netloc)) + validated_netloc = f"[{current_host}]" if ":" in current_host else current_host + if cp.port: + validated_netloc = f"{validated_netloc}:{cp.port}" + if os.environ.get(_DISABLE_DNS_PINNING_ENV) == "1": + # Enterprise proxies need the hostname in CONNECT for policy and TLS interception. + request_url = urlunparse(cp._replace(netloc = validated_netloc)) + else: + # Pin to the validated IP to prevent DNS rebinding. + ip_str = f"[{pinned_ip}]" if ":" in pinned_ip else pinned_ip + ip_netloc = f"{ip_str}:{cp.port}" if cp.port else ip_str + request_url = urlunparse(cp._replace(netloc = ip_netloc)) opener = urllib.request.build_opener( _NoRedirect, @@ -4209,11 +4215,11 @@ def _fetch_url_raw( headers = { "User-Agent": ua, - "Host": current_host, + "Host": validated_netloc, } if extra_headers: headers.update(extra_headers) - req = urllib.request.Request(pinned_url, headers = headers) + req = urllib.request.Request(request_url, headers = headers) try: # Cap the socket timeout at the time left on the overall deadline # so a single slow hop cannot outlast the whole fetch budget. diff --git a/studio/backend/run.py b/studio/backend/run.py index 90fd28670c..f1fc8c6062 100644 --- a/studio/backend/run.py +++ b/studio/backend/run.py @@ -1885,6 +1885,12 @@ def _build_arg_parser(): default = None, help = "Force server-side tools off for every request.", ) + parser.add_argument( + "--disable-dns-pinning", + action = "store_true", + help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens " + "DNS-rebinding protection; hostname and redirect validation remain enabled.", + ) parser.add_argument( "--parallel", "--n-parallel", @@ -1924,6 +1930,10 @@ if __name__ == "__main__": parser.error( "--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare" ) + if args.disable_dns_pinning: + os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1" + else: + os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0") kwargs = dict( host = args.host, diff --git a/studio/backend/tests/test_secure_tunnel_gate.py b/studio/backend/tests/test_secure_tunnel_gate.py index a8c0c2305f..b491134045 100644 --- a/studio/backend/tests/test_secure_tunnel_gate.py +++ b/studio/backend/tests/test_secure_tunnel_gate.py @@ -85,6 +85,14 @@ def test_arg_parser_secure_polarity_and_not_secure_alias(): assert parser.parse_args(["--not-secure", "--secure"]).secure is True +def test_arg_parser_dns_pinning_opt_out_defaults_off(): + import run + + parser = run._build_arg_parser() + assert parser.parse_args([]).disable_dns_pinning is False + assert parser.parse_args(["--disable-dns-pinning"]).disable_dns_pinning is True + + def test_run_server_accepts_enable_tools_kwarg(): import inspect diff --git a/studio/backend/tests/test_web_fetch_extraction.py b/studio/backend/tests/test_web_fetch_extraction.py index b794ee3e81..d4c3d123c3 100644 --- a/studio/backend/tests/test_web_fetch_extraction.py +++ b/studio/backend/tests/test_web_fetch_extraction.py @@ -15,6 +15,8 @@ from __future__ import annotations import sys from pathlib import Path +import pytest + _BACKEND_DIR = str(Path(__file__).resolve().parent.parent) if _BACKEND_DIR not in sys.path: sys.path.insert(0, _BACKEND_DIR) @@ -715,6 +717,61 @@ def test_fetch_url_raw_missing_content_type_reported_empty(monkeypatch): assert content_type == "" +@pytest.mark.parametrize( + "disable_dns_pinning,expected_url", + [ + (False, "https://203.0.113.7:8443/page?q=1"), + (True, "https://example.com:8443/page?q=1"), + ], +) +def test_fetch_url_raw_dns_pinning_proxy_opt_out(monkeypatch, disable_dns_pinning, expected_url): + import email + import urllib.request + + import core.inference.tools as tools_mod + + class _FakeResp: + headers = email.message_from_string("Content-Type: text/plain\n") + + def __init__(self): + self._body = b"ok" + + def read(self, n = -1): + body, self._body = self._body, b"" + return body + + requested = [] + + class _FakeOpener: + def open( + self, + req, + timeout = None, + ): + requested.append(req) + return _FakeResp() + + resolved = [] + + def resolve(host, port): + resolved.append((host, port)) + return True, "", "203.0.113.7" + + monkeypatch.setenv("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "1" if disable_dns_pinning else "0") + monkeypatch.setattr(tools_mod, "_validate_and_resolve_host", resolve) + monkeypatch.setattr(urllib.request, "build_opener", lambda *handlers: _FakeOpener()) + + err, body, _content_type = tools_mod._fetch_url_raw( + "https://user:secret@example.com:8443/page?q=1" + ) + + assert err is None + assert body == "ok" + assert resolved == [("example.com", 8443)] + assert [req.full_url for req in requested] == [expected_url] + assert requested[0].get_header("Host") == "example.com:8443" + + def test_fetch_page_text_missing_content_type_html_sniffed(monkeypatch): # A header-less server returning an HTML body must still be converted. def fake_fetch( diff --git a/tests/studio/test_cli_studio_defaults.py b/tests/studio/test_cli_studio_defaults.py index b03fa57603..a39956fab0 100644 --- a/tests/studio/test_cli_studio_defaults.py +++ b/tests/studio/test_cli_studio_defaults.py @@ -68,3 +68,10 @@ def test_studio_run_host_is_loopback(): f"`unsloth studio run` --host default must be '127.0.0.1' (loopback) " f"but got '{host_default}'." ) + + +def test_dns_pinning_opt_out_is_registered_safe_by_default(): + source = _STUDIO_CMD_PY.read_text() + for func_name in ("studio_default", "run"): + default = _find_typer_option_default(source, func_name, "--disable-dns-pinning") + assert default is False, f"{func_name} must keep DNS pinning enabled by default" diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 84c22aa4ac..c41963aab9 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -1285,6 +1285,12 @@ def studio_default( help = "Force server-side tools (web search, code execution) on or off for " "every request. Default: on for every bind, with the per-chat UI toggle honored.", ), + disable_dns_pinning: bool = typer.Option( + False, + "--disable-dns-pinning", + help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens " + "DNS-rebinding protection; hostname and redirect validation remain enabled.", + ), password: str = typer.Option( "", "--password", @@ -1354,6 +1360,15 @@ def studio_default( err = True, ) raise typer.Exit(2) + if disable_dns_pinning: + typer.echo( + "Error: --disable-dns-pinning on `unsloth studio` applies to the " + f"plain-server path only. For `unsloth studio {ctx.invoked_subcommand}`, " + f"put it after the subcommand: `unsloth studio {ctx.invoked_subcommand} " + "--disable-dns-pinning ...`", + err = True, + ) + raise typer.Exit(2) # Same for --api-only: dropping it here would silently serve the UI. if api_only: typer.echo( @@ -1398,6 +1413,10 @@ def studio_default( # default (plain-server path; the `run` subcommand has its own --verbose). if verbose: _enable_verbose_access_logs() + if disable_dns_pinning: + os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1" + else: + os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0") # Use the studio venv if present and not already in it. Resolve the child # launcher BEFORE the gate: a headless gate strips the seeded @@ -1739,6 +1758,13 @@ def run( "every request. Default: on for every bind." ), ), + disable_dns_pinning: bool = typer.Option( + False, + "--disable-dns-pinning", + rich_help_panel = _RUN_PANEL_TOOLS, + help = "Allow hostname-based web fetches for enterprise proxies. WARNING: weakens " + "DNS-rebinding protection; hostname and redirect validation remain enabled.", + ), tool_call_healing: Optional[bool] = typer.Option( None, "--enable-tool-call-healing/--disable-tool-call-healing", @@ -1944,6 +1970,10 @@ def run( _enable_verbose_access_logs() if not any(a in ("--verbose", "-v", "--log-verbose") for a in extra_llama_args): extra_llama_args.append("--log-verbose") + if disable_dns_pinning: + os.environ["UNSLOTH_STUDIO_DISABLE_DNS_PINNING"] = "1" + else: + os.environ.setdefault("UNSLOTH_STUDIO_DISABLE_DNS_PINNING", "0") # Promote legacy exact `-m`/`-hfr`/`-f` back into typer params; # clusters stay in extras. From 0a2a4e2e3252325eb2a80ecdde09b0dbe53857e8 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Sat, 25 Jul 2026 22:46:53 -0700 Subject: [PATCH 09/25] Settings: widen dialog to 960px and raise height to 680px (#7456) Also caps the height at the viewport instead of pinning it, so short viewports no longer get a clipped dialog. --- studio/frontend/src/features/settings/settings-dialog.tsx | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index 7a3a3c58f7..e3ef476470 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -249,9 +249,10 @@ export function SettingsDialog() { } }} className={cn( - // Cap at 880px but shrink to the viewport so it doesn't clip on - // iPad-portrait widths where a fixed width overflows. - "settings-surface !max-w-[min(880px,calc(100vw-2rem))] h-[560px] w-[min(880px,calc(100vw-2rem))] p-0 overflow-hidden", + // Cap at 960px but shrink to the viewport so it doesn't clip on + // iPad-portrait widths where a fixed width overflows. Height caps + // the same way so short viewports don't get a clipped dialog. + "settings-surface !max-w-[min(960px,calc(100vw-2rem))] h-[min(680px,calc(100dvh-2rem))] w-[min(960px,calc(100vw-2rem))] p-0 overflow-hidden", // Soft shadow, no outline ring. Pin --radius to the light value so // corner rounding matches in dark mode. "shadow-border rounded-xl ring-0 [--radius:1.1rem]", From 8dffde9611f6a137579ed385f224176428e4d344 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Sat, 25 Jul 2026 23:09:48 -0700 Subject: [PATCH 10/25] Sidebar: settings gear above the profile in the collapsed rail (#7458) The profile-row cog is hidden when the rail collapses, leaving no way to reach settings without opening the account menu. --- studio/frontend/src/components/app-sidebar.tsx | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index c95112c748..3cda2a8690 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -1855,6 +1855,18 @@ export function AppSidebar() { )} + {/* Collapsed rail has no room for the cog on the profile row, so it + sits above the avatar instead. */} + { + useSettingsDialogStore.getState().openDialog(); + closeMobileIfOpen(); + }} + /> From bac04ab577d33f27a94c5b75cc4667acfd8a7db9 Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Sat, 25 Jul 2026 23:54:48 -0700 Subject: [PATCH 11/25] Add drag and drop sources to the create project dialog (#7441) * feat(studio): add drag and drop sources to create project Files dropped on the create-project dialog upload to the new project's sources as soon as it exists, so a project can start with context instead of needing a second trip to the Sources tab. The sidebar and projects page dialogs now reuse NewProjectDialog rather than each keeping their own copy, and the OCR / caption ingest overrides move to a shared helper so every upload path sends the same settings. * fix(studio): harden project source drops Drops are not filtered by the `accept` attribute the way the picker is, so a folder or an image would stage and then fail server-side with a confusing per-file error. Unsupported entries are now refused up front with one message. Cancel bypassed the dialog's reset, so a discarded name and its staged files came back on reopen and uploaded into the next project created. Every close path now goes through one handler. Long filenames lost their extension in _sanitize_filename and were then rejected as an unsupported type; the stem is trimmed instead. Adds backend tests for the project scope, the sanitizer and path stripping. * fix(studio): address second review pass on source drops A drop landing on the panel while uploads run was not cancelled, because pointer-events-none took the panel out of hit testing and nothing else on the page cancels a file drop. The browser would navigate to the file and kill the uploads in flight. Drag defaults are now cancelled even while disabled, and the files are ignored instead. Name, size and mtime can match for two genuinely different files, so a skipped duplicate now says so rather than disappearing. A slow upload could resolve after the dialog unmounted and still navigate, pulling the user off the page they had moved to. Post-upload work is gated on the component still being mounted. * fix(studio): make source drops safe under StrictMode replay The mount sentinel was only cleared in effect cleanup, so StrictMode's setup/cleanup/setup replay left it false for good and every create in a dev build stopped short of closing the dialog or navigating. It is now set on setup as well. The pending-sources marker was consumed inside a useState initializer, which React replays, so the discarded pass ate the flag and the project opened on Chats. Reading is now a peek and the marker is dropped in an effect. Identical bytes under two names collapse to one document server-side, which looked like both files had been added. The upload loop now tracks returned document ids and says when files were merged. * fix(studio): guard the route and storage around staged uploads The sidebar's dialog lives in the root layout and never unmounts on a route change, so the mount check alone could not stop a slow upload from navigating the user back to the new project. The route is captured when create is pressed and compared afterwards, and callers get that answer so the sidebar can still move a chat while leaving the user where they are. Reading the vision-pass overrides went straight at localStorage, which throws outright where storage is blocked. That happened before the upload loop, so a project was created and every staged source was lost. It now falls back to the backend defaults, matching loadOptionalBool in the chat runtime store. --- studio/backend/routes/rag.py | 9 +- .../tests/test_rag_project_source_upload.py | 81 +++++ .../frontend/src/components/app-sidebar.tsx | 95 ++---- .../frontend/src/features/chat/chat-page.tsx | 13 +- .../chat/components/new-project-dialog.tsx | 141 ++++++-- .../src/features/chat/projects-page.tsx | 70 +--- .../components/project-source-dropzone.tsx | 304 ++++++++++++++++++ .../rag/components/use-rag-documents.ts | 21 +- .../rag/components/vision-overrides.ts | 35 ++ 9 files changed, 583 insertions(+), 186 deletions(-) create mode 100644 studio/backend/tests/test_rag_project_source_upload.py create mode 100644 studio/frontend/src/features/rag/components/project-source-dropzone.tsx create mode 100644 studio/frontend/src/features/rag/components/vision-overrides.ts diff --git a/studio/backend/routes/rag.py b/studio/backend/routes/rag.py index 392a4e0d02..ae65712146 100644 --- a/studio/backend/routes/rag.py +++ b/studio/backend/routes/rag.py @@ -47,7 +47,14 @@ _SAFE = re.compile(r"[^A-Za-z0-9._-]+") def _sanitize_filename(name: str) -> str: base = os.path.basename(name or "").strip() or "document" base = _SAFE.sub("_", base) - return base[:200] + if len(base) <= 200: + return base + # Trim the stem, not the extension: _save_upload gates on the extension, so + # a plain truncation would reject a long-named .txt as "unsupported". + stem, ext = os.path.splitext(base) + if not ext or len(ext) > 32: + return base[:200] + return stem[: 200 - len(ext)] + ext def _save_upload(file: UploadFile) -> tuple[str, str]: diff --git a/studio/backend/tests/test_rag_project_source_upload.py b/studio/backend/tests/test_rag_project_source_upload.py new file mode 100644 index 0000000000..fd20816b56 --- /dev/null +++ b/studio/backend/tests/test_rag_project_source_upload.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 + +"""Project sources upload: the path the create-project dialog drives.""" + +import os + +import pytest + +from core.rag import ingestion, store +from routes.rag import _sanitize_filename +from storage import rag_db + + +def _wait(job_id, timeout = 30.0): + import time + + deadline = time.time() + timeout + while time.time() < deadline: + status = ingestion.get_job_status(job_id) + if status and status["status"] in ("completed", "failed"): + return status + time.sleep(0.05) + raise AssertionError("ingestion did not finish in time") + + +def _ingest(project_id, filename, path): + return ingestion.start_ingestion( + store.project_scope(project_id), None, None, filename, path, project_id = project_id + ) + + +def test_project_document_persists_under_its_scope(rag_home, stub_embeddings, tmp_path): + path = tmp_path / "notes.txt" + path.write_text("alpha bravo charlie " * 50, encoding = "utf-8") + _, job_id = _ingest("P1", "notes.txt", str(path)) + assert _wait(job_id)["status"] == "completed" + + conn = rag_db.get_connection() + try: + docs = store.list_documents(conn, store.project_scope("P1")) + assert [d["filename"] for d in docs] == ["notes.txt"] + # Scoped: a sibling project cannot see it. + assert store.list_documents(conn, store.project_scope("P2")) == [] + assert store.search_lexical(conn, store.project_scope("P1"), "bravo", 5) + finally: + conn.close() + + +@pytest.mark.parametrize( + "raw", + [ + "x" * 300 + ".txt", + "y" * 512 + ".PDF", + "../" * 80 + "deep.md", + ], +) +def test_long_filenames_keep_their_extension(raw): + # _save_upload gates on the extension, so trimming it would reject the file. + out = _sanitize_filename(raw) + assert len(out) <= 200 + assert os.path.splitext(out)[1].lower() == os.path.splitext(raw)[1].lower() + + +@pytest.mark.parametrize( + "raw", + [ + "../../etc/passwd.txt", + "..\\..\\windows\\evil.txt", + "/absolute/notes.txt", + "C:\\Users\\me\\notes.txt", + ], +) +def test_sanitized_filenames_carry_no_path(raw): + out = _sanitize_filename(raw) + assert "/" not in out and "\\" not in out + + +@pytest.mark.parametrize("raw", ["." * 300, "noext" * 100, "a" * 100 + "." + "e" * 250]) +def test_sanitizer_degrades_safely(raw): + assert 0 < len(_sanitize_filename(raw)) <= 200 diff --git a/studio/frontend/src/components/app-sidebar.tsx b/studio/frontend/src/components/app-sidebar.tsx index 3cda2a8690..a31d9b6ced 100644 --- a/studio/frontend/src/components/app-sidebar.tsx +++ b/studio/frontend/src/components/app-sidebar.tsx @@ -105,7 +105,6 @@ import { archiveChatItem, ChatSearchDialog, clearNewChatDraft, - createChatProject, deleteChatProject, deleteChatItem, listStoredChatThreads, @@ -123,6 +122,7 @@ import { type ProjectRecord, type SidebarItem, } from "@/features/chat"; +import { NewProjectDialog } from "@/features/chat/components/new-project-dialog"; import { useAppearanceCustomStore, useSettingsDialogStore, @@ -696,7 +696,6 @@ export function AppSidebar() { }); }, [allChatItems, pendingRename]); const [creatingProject, setCreatingProject] = useState(false); - const [projectNameDraft, setProjectNameDraft] = useState(""); const [projectCreateMoveTarget, setProjectCreateMoveTarget] = useState(null); const renameTrimmed = renameDraft.trim(); @@ -849,28 +848,26 @@ export function AppSidebar() { } } - async function commitCreateProject() { - const name = projectNameDraft.trim(); - if (!name) return; + // "New project" from a chat's menu moves that chat in and stays put; + // otherwise open the project, unless a slow upload outlasted the route the + // user was on when they hit create. + async function afterCreateProject( + project: ProjectRecord, + { stayedOnRoute }: { stayedOnRoute: boolean }, + ) { const moveTarget = projectCreateMoveTarget; + setProjectCreateMoveTarget(null); + if (!moveTarget) { + if (stayedOnRoute) openProject(project.id); + return; + } try { - const project = await createChatProject(name); - if (moveTarget) { - await moveChatItemToProject(moveTarget, project.id); - if (activeThreadId === moveTarget.id) { - useChatRuntimeStore.getState().setActiveProjectId(project.id); - } - } - setCreatingProject(false); - setProjectNameDraft(""); - setProjectCreateMoveTarget(null); - if (moveTarget) { - return; - } else { - openProject(project.id); + await moveChatItemToProject(moveTarget, project.id); + if (activeThreadId === moveTarget.id) { + useChatRuntimeStore.getState().setActiveProjectId(project.id); } } catch (err) { - toast.error(moveTarget ? "Failed to create and move chat" : "Failed to create project", { + toast.error("Failed to move chat to the new project", { description: err instanceof Error ? err.message : undefined, }); } @@ -1050,7 +1047,6 @@ export function AppSidebar() { { setProjectCreateMoveTarget(item); - setProjectNameDraft(""); setCreatingProject(true); }} > @@ -1393,7 +1389,6 @@ export function AppSidebar() { onClick={(e) => { e.stopPropagation(); setProjectCreateMoveTarget(null); - setProjectNameDraft(""); setCreatingProject(true); }} className="sidebar-row-action group-hover/projects-item:opacity-100 group-hover/projects-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto group-data-[collapsible=icon]:hidden" @@ -2172,58 +2167,18 @@ export function AppSidebar() { - { setCreatingProject(open); - if (!open) { - setProjectNameDraft(""); - setProjectCreateMoveTarget(null); - } + if (!open) setProjectCreateMoveTarget(null); }} - > - - - - {projectCreateMoveTarget ? "Move to new project" : "New project"} - - - setProjectNameDraft(event.target.value)} - onKeyDown={(event) => { - if (event.key === "Enter") { - event.preventDefault(); - void commitCreateProject(); - } - }} - autoFocus - maxLength={120} - placeholder="Project name" - aria-label="Project name" - className="focus-visible:border-input focus-visible:ring-0" - /> - - - - - - + title={ + projectCreateMoveTarget ? "Move to new project" : "Create project" + } + submitLabel={projectCreateMoveTarget ? "Create and move" : "Create project"} + onCreated={afterCreateProject} + /> ); } diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index a439a91239..7e0544e2d1 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -185,6 +185,10 @@ import { listStoredChatThreads, } from "./utils/chat-history-storage"; import { isAssistantLocalThreadId } from "./utils/thread-ids"; +import { + consumeProjectSourcesPending, + hasProjectSourcesPending, +} from "@/features/rag/components/project-source-dropzone"; const ProjectSourcesPanel = lazy(() => @@ -998,7 +1002,14 @@ function ProjectLanding({ const active = useChatActive(); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const initialActiveThreadRef = useRef(null); - const [projectTab, setProjectTab] = useState<"chats" | "sources">("chats"); + // Land on Sources when the project was just created with dropped files. + const [projectTab, setProjectTab] = useState<"chats" | "sources">(() => + hasProjectSourcesPending(projectId) ? "sources" : "chats", + ); + // Drop the marker once committed: React may replay the initializer above. + useEffect(() => { + consumeProjectSourcesPending(projectId); + }, [projectId]); const [pendingNewThreadId, setPendingNewThreadId] = useState( null, ); diff --git a/studio/frontend/src/features/chat/components/new-project-dialog.tsx b/studio/frontend/src/features/chat/components/new-project-dialog.tsx index 880129ac6c..6aca3d36c3 100644 --- a/studio/frontend/src/features/chat/components/new-project-dialog.tsx +++ b/studio/frontend/src/features/chat/components/new-project-dialog.tsx @@ -2,7 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { useNavigate } from "@tanstack/react-router"; -import { useState } from "react"; +import { useEffect, useRef, useState } from "react"; import { Button } from "@/components/ui/button"; import { @@ -12,31 +12,92 @@ import { DialogHeader, DialogTitle, } from "@/components/ui/dialog"; -import { Input } from "@/components/ui/input"; +import { + ProjectSourceDropzone, + type StagedSource, + uploadStagedSources, +} from "@/features/rag/components/project-source-dropzone"; import { toast } from "@/lib/toast"; +import { Folder02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; import { createChatProject } from "../hooks/use-chat-projects"; import { useChatRuntimeStore } from "../stores/chat-runtime-store"; +import type { ProjectRecord } from "../types"; -// Create-project dialog usable from the composer + menu. Creating opens the new -// project straight away rather than dropping the user on the projects list. +function currentRoute(): string { + if (typeof window === "undefined") return ""; + return window.location.pathname + window.location.search; +} + +// Create-project dialog for the composer, sidebar, and projects page. Creating +// opens the new project; `onCreated` overrides that for callers with their own +// follow-up (the sidebar's "move this chat to a new project"). export function NewProjectDialog({ open, onOpenChange, + title = "Create project", + submitLabel = "Create project", + onCreated, }: { open: boolean; onOpenChange: (open: boolean) => void; + title?: string; + submitLabel?: string; + onCreated?: ( + project: ProjectRecord, + context: { stayedOnRoute: boolean }, + ) => void | Promise; }) { const navigate = useNavigate(); const [name, setName] = useState(""); + const [staged, setStaged] = useState([]); + const [busy, setBusy] = useState(false); + // Uploads outlive this component, so a slow one must not yank the user to the + // new project after they have navigated away. + const mounted = useRef(true); + useEffect(() => { + // Set on setup, not just cleared on cleanup: StrictMode replays + // setup/cleanup/setup, which would otherwise leave this false forever. + mounted.current = true; + return () => { + mounted.current = false; + }; + }, []); + + function reset() { + setName(""); + setStaged([]); + } + + // Every close path routes through here: callers keep this mounted, so a draft + // left behind would resurface (and upload) on the next project. + function close() { + if (busy) return; + reset(); + onOpenChange(false); + } async function commitCreate() { const trimmed = name.trim(); - if (!trimmed) return; + if (!trimmed || busy) return; + setBusy(true); + // Sidebar callers keep this mounted across routes, so unmounting alone + // cannot tell whether the user has moved on during a slow upload. + const origin = currentRoute(); try { const project = await createChatProject(trimmed); + // Upload before closing so the Sources panel lists them on first fetch. + await uploadStagedSources(project.id, staged); + if (!mounted.current) return; + const stayedOnRoute = currentRoute() === origin; onOpenChange(false); - setName(""); + reset(); + if (onCreated) { + await onCreated(project, { stayedOnRoute }); + return; + } + if (!stayedOnRoute) return; const runtime = useChatRuntimeStore.getState(); runtime.setActiveThreadId(null); runtime.setActiveProjectId(project.id); @@ -45,6 +106,8 @@ export function NewProjectDialog({ toast.error("Failed to create project", { description: err instanceof Error ? err.message : undefined, }); + } finally { + setBusy(false); } } @@ -52,43 +115,59 @@ export function NewProjectDialog({ { - if (!next) setName(""); - onOpenChange(next); + if (next) { + onOpenChange(true); + return; + } + close(); }} > - + - New project + {title} - setName(e.target.value)} - onKeyDown={(e) => { - if (e.key === "Enter") { - e.preventDefault(); - void commitCreate(); - } - }} - autoFocus={true} - maxLength={120} - placeholder="Project name" - aria-label="Project name" - className="focus-visible:border-input focus-visible:ring-0" + {/* Name field: folder glyph in its own cell, divided from the input. */} +
+ + + +
+ - diff --git a/studio/frontend/src/features/chat/projects-page.tsx b/studio/frontend/src/features/chat/projects-page.tsx index e20e517787..192c4e2331 100644 --- a/studio/frontend/src/features/chat/projects-page.tsx +++ b/studio/frontend/src/features/chat/projects-page.tsx @@ -34,7 +34,6 @@ import { isTauri } from "@/lib/api-base"; import { isDownloadCancelled, pickNativeChatImport } from "@/lib/native-files"; import { toast } from "@/lib/toast"; import { - createChatProject, deleteChatProject, renameChatProject, useChatProjects, @@ -42,6 +41,7 @@ import { usePinnedProjectsStore, type ProjectRecord, } from "@/features/chat"; +import { NewProjectDialog } from "./components/new-project-dialog"; import { Delete02Icon, Download01Icon, @@ -124,7 +124,6 @@ export function ProjectsPage() { ); const [creating, setCreating] = useState(false); - const [nameDraft, setNameDraft] = useState(""); const [renaming, setRenaming] = useState(null); const [renameDraft, setRenameDraft] = useState(""); const [deleting, setDeleting] = useState(null); @@ -258,21 +257,6 @@ export function ProjectsPage() { navigate({ to: "/chat", search: { project: projectId } }); } - async function commitCreate() { - const name = nameDraft.trim(); - if (!name) return; - try { - const project = await createChatProject(name); - setCreating(false); - setNameDraft(""); - openProject(project.id); - } catch (err) { - toast.error("Failed to create project", { - description: err instanceof Error ? err.message : undefined, - }); - } - } - async function commitRename() { const target = renaming; const name = renameDraft.trim(); @@ -469,14 +453,7 @@ export function ProjectsPage() { - + @@ -511,10 +488,7 @@ export function ProjectsPage() { - - - - + {/* Create project (name + drag-and-drop sources) */} + {/* Rename project */} = 1024 && unit < units.length - 1) { + value /= 1024; + unit += 1; + } + const shown = + value >= 10 || unit === 0 + ? String(Math.round(value)) + : value.toFixed(1).replace(/\.0$/, ""); + return `${shown} ${units[unit]}`; +} + +const ACCEPTED_EXTS = new Set( + RAG_UPLOAD_ACCEPT.split(",").map((ext) => ext.trim().toLowerCase()), +); + +// `accept` only filters the picker, so a drop can carry anything. A folder +// arrives as an extension-less entry, which this rejects along with the types +// the backend would 400 on. +function isSupported(file: File): boolean { + const dot = file.name.lastIndexOf("."); + if (dot <= 0) return false; + return ACCEPTED_EXTS.has(file.name.slice(dot).toLowerCase()); +} + +/** Merge a selection into the staged list. Returns the names it would not take, + * so the caller can say so once instead of dropping them silently. */ +function addStagedSources( + staged: StagedSource[], + incoming: FileList | File[], +): { next: StagedSource[]; unsupported: string[]; duplicates: string[] } { + const seen = new Set(staged.map((entry) => fileSignature(entry.file))); + const next = [...staged]; + const unsupported: string[] = []; + const duplicates: string[] = []; + for (const file of Array.from(incoming)) { + if (!isSupported(file)) { + unsupported.push(file.name); + continue; + } + const signature = fileSignature(file); + if (seen.has(signature)) { + duplicates.push(file.name); + continue; + } + seen.add(signature); + next.push({ + id: `staged_${Math.random().toString(36).slice(2)}`, + file, + }); + } + return { next, unsupported, duplicates }; +} + +// Projects created with staged files, so the landing can open on Sources. +const projectsWithPendingSources = new Set(); + +function markProjectSourcesPending(projectId: string): void { + projectsWithPendingSources.add(projectId); +} + +/** Whether this project was just created with staged sources. Read-only, so it + * is safe in a render pass that React may replay. */ +export function hasProjectSourcesPending(projectId: string): boolean { + return projectsWithPendingSources.has(projectId); +} + +/** Drop the marker once the landing has committed. */ +export function consumeProjectSourcesPending(projectId: string): void { + projectsWithPendingSources.delete(projectId); +} + +/** Upload staged files to a new project. Indexing runs in the background; a + * per-file failure toasts and never blocks project creation. */ +export async function uploadStagedSources( + projectId: string, + staged: StagedSource[], +): Promise { + if (staged.length === 0) return; + invalidateProjectSources(projectId); + markProjectSourcesPending(projectId); + const { ocr, caption } = resolveVisionOverrides(); + const documentIds = new Set(); + const merged: string[] = []; + for (const { file } of staged) { + try { + const result = await uploadProjectDocument(projectId, file, ocr, caption); + // Same bytes under another name: the backend hashes content, so this is + // the document already uploaded. Say so rather than imply a new source. + if (documentIds.has(result.documentId)) merged.push(file.name); + else documentIds.add(result.documentId); + } catch (error) { + toast.error(`Couldn't upload ${file.name}`, { + description: error instanceof Error ? error.message : String(error), + }); + } + } + if (merged.length > 0) { + toast.info( + merged.length === 1 + ? `${merged[0]} matched a file already added` + : `${merged.length} files matched files already added`, + { description: "Identical contents are stored once." }, + ); + } + invalidateProjectSources(projectId); +} + +/** Create-project drop area: stages files until the project exists. */ +export function ProjectSourceDropzone({ + staged, + onChange, + disabled = false, +}: { + staged: StagedSource[]; + onChange: (next: StagedSource[]) => void; + disabled?: boolean; +}) { + const inputRef = useRef(null); + // Count enter/leave pairs: children fire dragleave on the parent. + const dragDepth = useRef(0); + const [dragging, setDragging] = useState(false); + + const addFiles = useCallback( + (files: FileList | File[]) => { + const { next, unsupported, duplicates } = addStagedSources(staged, files); + if (next.length !== staged.length) onChange(next); + if (unsupported.length > 0) { + toast.info( + unsupported.length === 1 + ? `Can't add ${unsupported[0]}` + : `Can't add ${unsupported.length} files`, + { description: `Supported types: ${RAG_UPLOAD_ACCEPT}` }, + ); + } + // Name, size and mtime can in principle match for two different files, so + // never drop one without saying so. + if (duplicates.length > 0) { + toast.info( + duplicates.length === 1 + ? `${duplicates[0]} is already added` + : `${duplicates.length} files were already added`, + ); + } + }, + [staged, onChange], + ); + + const endDrag = useCallback(() => { + dragDepth.current = 0; + setDragging(false); + }, []); + + return ( +
+

Sources

+ {/* Panel is the drop target; the inner button owns the click so staged + rows can carry their own remove buttons. */} +
{ + e.preventDefault(); + if (disabled) return; + dragDepth.current += 1; + setDragging(true); + }} + onDragOver={(e) => { + e.preventDefault(); + if (disabled) return; + e.dataTransfer.dropEffect = "copy"; + }} + onDragLeave={() => { + dragDepth.current = Math.max(0, dragDepth.current - 1); + if (dragDepth.current === 0) setDragging(false); + }} + onDrop={(e) => { + e.preventDefault(); + if (disabled) return; + endDrag(); + addFiles(Array.from(e.dataTransfer.files ?? [])); + }} + className={cn( + "rounded-[22px] border border-border transition-colors dark:border-white/10", + dragging && "border-primary/60 bg-primary/5", + disabled && "opacity-60", + )} + > + { + const files = Array.from(e.target.files ?? []); + e.target.value = ""; + addFiles(files); + }} + /> + {staged.length === 0 ? ( + + ) : ( +
+
    + {staged.map((entry) => ( +
  • + + + {entry.file.name} + + + {formatSize(entry.file.size)} + + +
    • + ))} +
+ +
+ )} +
+
+ ); +} diff --git a/studio/frontend/src/features/rag/components/use-rag-documents.ts b/studio/frontend/src/features/rag/components/use-rag-documents.ts index 8d6433d8c3..bdab7b0518 100644 --- a/studio/frontend/src/features/rag/components/use-rag-documents.ts +++ b/studio/frontend/src/features/rag/components/use-rag-documents.ts @@ -1,13 +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 -import { useCallback, useEffect, useRef, useState } from "react"; -import { - CHAT_RAG_CAPTION_KEY, - CHAT_RAG_OCR_KEY, - useChatRuntimeStore, -} from "@/features/chat"; import { toast } from "@/lib/toast"; +import { useCallback, useEffect, useRef, useState } from "react"; import { deleteDocument, getJob, @@ -17,6 +12,7 @@ import { uploadThreadDocument, } from "../api/rag-api"; import type { DocumentStatus, RagDocument } from "../types/rag"; +import { resolveVisionOverrides } from "./vision-overrides"; export interface TrackedDocument extends RagDocument { progress?: number | null; @@ -263,18 +259,7 @@ export function useRagDocuments( tempId: string, ) => { try { - // Send vision-pass overrides only after the user has explicitly set them; - // otherwise backend env defaults own the ingest policy. - const state = useChatRuntimeStore.getState(); - const hasLocal = (key: string) => - typeof window !== "undefined" && - window.localStorage.getItem(key) !== null; - const ocr = hasLocal(CHAT_RAG_OCR_KEY) - ? state.ragOcrScanned - : undefined; - const caption = hasLocal(CHAT_RAG_CAPTION_KEY) - ? state.ragCaptionFigures - : undefined; + const { ocr, caption } = resolveVisionOverrides(); const result = activeScope.type === "kb" ? await uploadKnowledgeBaseDocument( diff --git a/studio/frontend/src/features/rag/components/vision-overrides.ts b/studio/frontend/src/features/rag/components/vision-overrides.ts new file mode 100644 index 0000000000..674484970b --- /dev/null +++ b/studio/frontend/src/features/rag/components/vision-overrides.ts @@ -0,0 +1,35 @@ +// 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 { + CHAT_RAG_CAPTION_KEY, + CHAT_RAG_OCR_KEY, + useChatRuntimeStore, +} from "@/features/chat"; + +function hasLocal(key: string): boolean { + if (typeof window === "undefined") return false; + try { + return window.localStorage.getItem(key) !== null; + } catch { + // Storage can be blocked outright (sandboxed context). These overrides are + // optional, so fall back to the backend defaults rather than failing the + // upload that asked for them. + return false; + } +} + +/** Ingest-time vision-pass overrides, sent only once the user has set them; + * otherwise backend env defaults own the policy. Shared by every upload path. */ +export function resolveVisionOverrides(): { + ocr: boolean | undefined; + caption: boolean | undefined; +} { + const state = useChatRuntimeStore.getState(); + return { + ocr: hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined, + caption: hasLocal(CHAT_RAG_CAPTION_KEY) + ? state.ragCaptionFigures + : undefined, + }; +} From 671d6dbf6902f0355496e5d42219daad5d2f06fb Mon Sep 17 00:00:00 2001 From: Michael Han <107991372+shimmyshimmer@users.noreply.github.com> Date: Sun, 26 Jul 2026 00:01:22 -0700 Subject: [PATCH 12/25] Settings: match dialog fills to the app shell surfaces (#7457) * Settings: match dialog fills to the app shell surfaces Tabs use the sidebar fill and the content pane uses the page fill, so both track the active palette in light and dark. * Pair the tab column fill with the sidebar foreground Custom themes set --foreground but not --sidebar, so search result rows could land white on white. Track the sidebar token instead. --- studio/frontend/src/features/settings/settings-dialog.tsx | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/studio/frontend/src/features/settings/settings-dialog.tsx b/studio/frontend/src/features/settings/settings-dialog.tsx index e3ef476470..0ba59f7095 100644 --- a/studio/frontend/src/features/settings/settings-dialog.tsx +++ b/studio/frontend/src/features/settings/settings-dialog.tsx @@ -267,7 +267,9 @@ export function SettingsDialog() { {/* Keep tab content from expanding the dialog grid. */}
-