studio: follow-up polish from autoload + launcher PR review rounds

autoLoad cascade hardening (extends #5578):

- Singleflight wrapper around autoLoadSmallestModel so two send-button
  clicks in flight share one cascade instead of fanning out a fresh
  /load budget per click.
- /validate POSTs were previously uncapped; a cache of N trust-blocked
  repos still produced N /validate calls. Adds MAX_AUTO_VALIDATE_ATTEMPTS
  = 10 with a counter inside canAutoLoad.
- abortSignal from the adapter's run() now threads into the cascade and
  is checked on entry plus at each loop boundary, so a tab close or
  user abort mid-cascade stops further POSTs.

Style consistency (extends #5577):

- _release_self_exe_lock_windows / _restore_self_exe_lock_windows now
  emit via typer.echo to match the rest of the file's diagnostic
  surface.
This commit is contained in:
Daniel Han 2026-05-19 13:04:31 +00:00
commit faf35ed2c0
2 changed files with 39 additions and 5 deletions

View file

@ -439,8 +439,32 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise<void> {
*/
// Cap cascade so broken cached repos can't spam /api/inference/load.
const MAX_AUTO_LOAD_ATTEMPTS = 3;
// /validate is cheaper but a long trust-blocked cache still produces one POST
// per repo; cap that path separately so the cascade is fully bounded.
const MAX_AUTO_VALIDATE_ATTEMPTS = 10;
async function autoLoadSmallestModel(): Promise<{
// Singleflight handle: concurrent sends arriving while a cascade is in
// flight share its result instead of fanning out another /load budget.
let autoLoadInflight: Promise<{
loaded: boolean;
blockedByTrustRemoteCode: boolean;
}> | null = null;
async function autoLoadSmallestModel(abortSignal?: AbortSignal): Promise<{
loaded: boolean;
blockedByTrustRemoteCode: boolean;
}> {
if (autoLoadInflight) return autoLoadInflight;
const work = runAutoLoadCascade(abortSignal);
autoLoadInflight = work;
try {
return await work;
} finally {
autoLoadInflight = null;
}
}
async function runAutoLoadCascade(abortSignal?: AbortSignal): Promise<{
loaded: boolean;
blockedByTrustRemoteCode: boolean;
}> {
@ -455,6 +479,7 @@ async function autoLoadSmallestModel(): Promise<{
let blockedByTrustRemoteCode = false;
let hadNonTrustFailure = false;
let loadAttempts = 0;
let validateAttempts = 0;
async function canAutoLoad(payload: {
model_path: string;
@ -462,6 +487,8 @@ async function autoLoadSmallestModel(): Promise<{
is_lora: boolean;
gguf_variant?: string | null;
}): Promise<boolean> {
if (validateAttempts >= MAX_AUTO_VALIDATE_ATTEMPTS) return false;
validateAttempts += 1;
const validation = await validateModel({
...payload,
hf_token: hfToken,
@ -480,11 +507,17 @@ async function autoLoadSmallestModel(): Promise<{
listCachedModels().catch(() => []),
]);
if (abortSignal?.aborted) {
toast.dismiss(toastId);
return { loaded: false, blockedByTrustRemoteCode: false };
}
// Try GGUF first: pick the repo with the smallest total size,
// then pick its smallest downloaded variant.
if (ggufRepos.length > 0) {
const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes);
for (const repo of sorted) {
if (abortSignal?.aborted) break;
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
try {
const variants = await listGgufVariants(repo.repo_id);
@ -566,6 +599,7 @@ async function autoLoadSmallestModel(): Promise<{
if (modelRepos.length > 0) {
const sorted = [...modelRepos].sort((a, b) => a.size_bytes - b.size_bytes);
for (const repo of sorted) {
if (abortSignal?.aborted) break;
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
try {
if (
@ -626,7 +660,7 @@ async function autoLoadSmallestModel(): Promise<{
// Cap also gates the default download so the total /api/inference/load
// budget across cached + fallback is MAX_AUTO_LOAD_ATTEMPTS, not +1.
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) {
if (abortSignal?.aborted || loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) {
toast.dismiss(toastId);
return {
loaded: false,
@ -736,7 +770,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
if (!useChatRuntimeStore.getState().params.checkpoint) {
// Auto-load the smallest downloaded model
const { loaded, blockedByTrustRemoteCode } =
await autoLoadSmallestModel();
await autoLoadSmallestModel(abortSignal);
if (!loaded) {
toast.error(
blockedByTrustRemoteCode

View file

@ -1315,7 +1315,7 @@ def _release_self_exe_lock_windows() -> None:
os.replace(exe, stale)
except OSError as e:
# Not fatal; setup.ps1 retries from a sibling process.
print(f"[update] could not rename {exe.name} -> {stale.name}: {e}")
typer.echo(f"[update] could not rename {exe.name} -> {stale.name}: {e}")
def _restore_self_exe_lock_windows() -> None:
@ -1341,7 +1341,7 @@ def _restore_self_exe_lock_windows() -> None:
try:
os.replace(stale, exe)
except OSError as e:
print(f"[update] could not restore {stale.name} -> {exe.name}: {e}")
typer.echo(f"[update] could not restore {stale.name} -> {exe.name}: {e}")
def _cleanup_self_exe_lock_windows() -> None: