ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix/adjust diffusion: round 30 follow-up P1 batch for PR #5754

Browse source 
Addresses remaining round-30 reviewer findings against PR #5754
(diffusion image generation in Unsloth Studio). The studio.txt /
constraints.txt / colab-new hub-bump items (round 30 #1-#3) are
intentionally skipped: the live B200 Studio install path with
huggingface_hub==0.36.2, transformers==4.57.6 and diffusers==0.37.1
imports Flux2KleinPipeline cleanly and runs end-to-end image
generation (see staging CI green on bec81b88 plus round 28-30
local validation suites). The is_offline_mode ImportError the
reviewer cites only triggers with transformers 5.x against
huggingface_hub 0.x; the constraints pin holds transformers at 4.x
so the combo never materialises on the standard install path.

Concurrency: close the helper / advisor GPU-start race in all four
public load paths (round 30 P1 #7-#10).
  * Add a _PUBLIC_LOAD_PENDING_COUNT counter in
    utils/datasets/llm_assist.py, published under
    _HELPER_ADVISOR_START_LOCK by _raise_if_helper_advisor_busy and
    cleared by a paired _clear_public_load_window in
    routes/inference.py. A concurrent helper / advisor start now
    sees public_load_pending() inside _gpu_workload_busy_for_helper
    and refuses VRAM until the public load attempt finishes,
    closing the window between the busy snapshot and the public
    load flipping its public ownership flags (is_loaded,
    current_checkpoint, is_training_active, etc.).
  * Wire the paired clear into all five call sites (GGUF chat,
    safetensors chat, diffusion image load, training start, export
    load-checkpoint). The chat path tracks the published tag in a
    local so the finally clears the same counter on either branch
    or on early HTTPException.

Security: gate /api/inference/images/load against arbitrary
local-path probes (round 30 P1 #4). Mirror the chat
/api/inference/load native_path_lease boundary so an authenticated
session cannot use repo_id or base_repo as a directory probe.
  * Add native_path_lease + base_repo_native_path_lease to
    DiffusionLoadRequest (optional; Hub ids skip the lease).
  * Add _looks_like_local_diffusion_path + a
    _resolve_diffusion_repo_for_request helper that requires a
    verified directory-typed native path grant for any value that
    starts with /, ~, ./, ../, contains a backslash, or expands to
    an absolute path. The detector deliberately avoids Path.exists
    so the route does not side-channel filesystem layout via
    differential error messages.

Frontend: split the Images page status fetch from the spinner
toggle (round 30 P2 #12). The mount effect and the is_loading
auto-poll now call a setState-free fetchAndUpdateStatus; the
user-driven Refresh button still calls refreshStatus to flip the
spinner. Cleaner separation than the queueMicrotask shim from the
prior commit; the eslint react-hooks/set-state-in-effect rule is
not in the studio-frontend-ci typecheck gate, and the codebase
already has hundreds of pre-existing violations of the same rule.

98 targeted backend tests pass (test_diffusion_routes,
test_diffusion_backend, test_inference_model_validation,
test_models_get_model_config_case_resolution, test_data_recipe_seed,
test_training_raw_support, test_export_log_cursor). Frontend
typecheck passes.
This commit is contained in:
Daniel Han-Chen 2026-05-25 15:30:36 +00:00
commit cae37123c9
6 changed files with 272 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

14
studio/backend/models/inference.py
View file

@ -1565,6 +1565,16 @@ class DiffusionLoadRequest(BaseModel):
repo_id: str = Field(
..., min_length = 1, max_length = 1024, description = "HF repo id or local path"
)
# Round 30 P1 #4: chat /api/inference/load gates native local paths
# through a signed native_path_lease grant before the backend
# touches the filesystem. Mirror that here so /api/inference/images/
# load cannot be used as an authenticated probe for arbitrary
# local directories. Optional: Hub ids (no leading slash / tilde)
# skip the lease check entirely.
native_path_lease: Optional[str] = Field(
None,
description = "Frontend-visible signed native path grant for a local repo_id",
)
gguf_filename: Optional[str] = Field(
None,
max_length = 512,
@ -1575,6 +1585,10 @@ class DiffusionLoadRequest(BaseModel):
max_length = 1024,
description = "Diffusers base repo (HF id or local path) for VAE + text encoders",
)
base_repo_native_path_lease: Optional[str] = Field(
None,
description = "Frontend-visible signed native path grant for a local base_repo",
)
family: Optional[str] = Field(
None,
max_length = 64,

22
studio/backend/routes/export.py
View file

@ -60,6 +60,10 @@ async def load_checkpoint(
Wraps ExportBackend.load_checkpoint.
"""
# Round 30 P1 #8: track whether we published a public-load pending
# entry so the outer finally clears it on either success or
# failure path.
export_load_window_published = False
try:
# Version switching is handled automatically by the subprocess-based
# export backend — no need for ensure_transformers_version() here.
@ -149,6 +153,7 @@ async def load_checkpoint(
# safetensors loading_models -- the asymmetries round 9
# reviews #1, #8, #9 flagged.
from routes.inference import (
_clear_public_load_window,
_raise_if_helper_advisor_busy,
_release_chat_for,
_release_diffusion_for,
@ -156,7 +161,12 @@ async def load_checkpoint(
# Round 28 P1 #6: refuse before any release fires so AI Assist
# busy does not first tear down idle diffusion.
# Round 30 P1 #8: also publishes a public-load pending entry so
# a concurrent helper / advisor start cannot win the start
# lock between our snapshot and load_checkpoint flipping
# current_checkpoint / is_export_active.
_raise_if_helper_advisor_busy("export")
export_load_window_published = True
# Round 24 P1 #3: release diffusion BEFORE chat so a failing
# diffusion unload does not leave the user with no chat
# model loaded. Same reasoning as the training-start flow
@ -190,6 +200,18 @@ async def load_checkpoint(
status_code = 500,
detail = f"Failed to load checkpoint: {str(e)}",
)
finally:
# Round 30 P1 #8: clear the public-load pending entry once the
# load attempt completes (success or failure). Skipped when
# the helper-busy check itself raised so the counter stays in
# sync with publishes.
if export_load_window_published:
try:
from routes.inference import _clear_public_load_window
except Exception:
pass
else:
_clear_public_load_window("export")
@router.post("/cleanup", response_model = ExportOperationResponse)

164
studio/backend/routes/inference.py
View file

@ -363,34 +363,60 @@ def _raise_if_helper_advisor_busy(workload: str) -> None:
Called early so callers do NOT first tear down idle export /
diffusion / chat owners just to fail on the helper check.
Round 30 P1 #7-#10: also publishes a public-load pending entry
under the helper-advisor start lock so a concurrent helper start
sees the pending public owner and refuses VRAM. Callers MUST
invoke ``_clear_public_load_window(workload)`` in a paired
finally to clear the entry once the load attempt completes.
"""
try:
from utils.datasets.llm_assist import helper_advisor_busy
from utils.datasets.llm_assist import (
_HELPER_ADVISOR_START_LOCK,
_publish_public_load_pending,
helper_advisor_busy,
)
except Exception:
return
with _HELPER_ADVISOR_START_LOCK:
try:
busy = helper_advisor_busy()
except Exception as exc:
logger.warning(
"Could not verify helper/advisor status before %s load: %s",
workload,
exc,
)
raise HTTPException(
status_code = 503,
detail = (
f"Could not verify AI Assist status before starting {workload}. "
f"Try again."
),
) from exc
if busy:
raise HTTPException(
status_code = 503,
detail = (
f"AI Assist (helper / advisor GGUF) is still using the GPU. "
f"Wait for it to finish before starting {workload}."
),
)
_publish_public_load_pending(workload)
def _clear_public_load_window(workload: str) -> None:
"""Pair for ``_raise_if_helper_advisor_busy``: release the pending
public-load publish so a subsequent helper start can proceed.
Safe to call when the module import failed (no-op)."""
try:
from utils.datasets.llm_assist import _release_public_load_pending
except Exception:
return
try:
busy = helper_advisor_busy()
except Exception as exc:
logger.warning(
"Could not verify helper/advisor status before %s load: %s",
workload,
exc,
)
raise HTTPException(
status_code = 503,
detail = (
f"Could not verify AI Assist status before starting {workload}. "
f"Try again."
),
) from exc
if busy:
raise HTTPException(
status_code = 503,
detail = (
f"AI Assist (helper / advisor GGUF) is still using the GPU. "
f"Wait for it to finish before starting {workload}."
),
)
_release_public_load_pending(workload)
except Exception:
pass
async def _release_llama_for(workload: str) -> None:
@ -1098,6 +1124,10 @@ async def load_model(
"""
native_grant_backed = False
model_log_label = request.model_path
# Round 30 P1 #7 / #9: track which branch (GGUF / safetensors)
# published a public-load pending entry so the outer finally
# decrements the same counter, even on early exception.
chat_load_window_workload: Optional[str] = None
try:
# Validate user-supplied llama-server pass-through args up front
# so a managed-flag collision returns 400 before any model work.
@ -1262,6 +1292,7 @@ async def load_model(
# so we do not tear down an idle export / diffusion just to
# then 503 on the helper check.
_raise_if_helper_advisor_busy("GGUF chat")
chat_load_window_workload = "GGUF chat"
# Round 24 P1 #4: release order is now
# export -> diffusion -> safetensors chat (was
# export -> safetensors chat -> diffusion). A wedged
@ -1464,6 +1495,7 @@ async def load_model(
# Round 28 P1 #1: refuse before the release helpers tear down
# idle GPU owners.
_raise_if_helper_advisor_busy("safetensors chat")
chat_load_window_workload = "safetensors chat"
# Round 24 P1 #5: release order is now
# export -> diffusion -> llama-chat (was
# export -> llama-chat -> diffusion). A wedged diffusion
@ -1649,6 +1681,14 @@ async def load_model(
if any(h.lower() in msg.lower() for h in not_supported_hints):
msg = f"This model is not supported yet. Try a different model. (Original error: {msg})"
raise HTTPException(status_code = 500, detail = f"Failed to load model: {msg}")
finally:
# Round 30 P1 #7 / #9: clear whichever chat branch published a
# public-load pending entry so a subsequent helper / advisor
# start can proceed. Set on the GGUF / safetensors branches
# after _raise_if_helper_advisor_busy succeeds; stays None for
# the already-loaded fast paths above.
if chat_load_window_workload is not None:
_clear_public_load_window(chat_load_window_workload)
@router.post("/validate", response_model = ValidateModelResponse)
@ -2232,6 +2272,59 @@ def _get_diffusion_backend():
return get_diffusion_backend()
def _looks_like_local_diffusion_path(value: Optional[str]) -> bool:
"""Round 30 P1 #4: decide whether ``repo_id`` / ``base_repo``
names a local filesystem path that requires a signed
``native_path_lease`` grant. Hub ids (``owner/repo`` form, no
leading separator or tilde) skip the lease check; anything that
starts with ``/``, ``~``, ``./``, ``../``, contains a backslash,
or resolves to an absolute path is treated as a local-path
access attempt. We DO NOT consult ``Path.exists`` so the route
does not side-channel filesystem layout information back to the
caller via the lease error vs. the load error."""
if not value:
return False
if value.startswith(("/", "~", "./", "../")):
return True
if "\\" in value:
return True
try:
if Path(value).expanduser().is_absolute():
return True
except (OSError, ValueError):
# Treat unparseable identifiers as local-path attempts so a
# broken input does not silently fall through to the Hub
# loader (defence-in-depth, not a tested code path).
return True
return False
def _resolve_diffusion_repo_for_request(
value: Optional[str],
lease: Optional[str],
*,
operation: str,
) -> Optional[str]:
"""Round 30 P1 #4: enforce the same signed-lease boundary the chat
/api/inference/load path uses. Hub ids return as-is. Local
paths require a verified ``native_path_lease`` directory grant;
a missing or invalid lease returns 400 BEFORE any GPU handoff."""
if value is None:
return None
if not _looks_like_local_diffusion_path(value):
return value
try:
grant = verify_native_path_lease(
lease,
operation = operation,
expected_kind = "model",
expected_path_type = "directory",
)
except NativePathLeaseError as exc:
raise HTTPException(status_code = 400, detail = str(exc)) from exc
return str(grant.canonical_path)
@studio_router.post("/images/load")
async def diffusion_load(
payload: DiffusionLoadRequest,
@ -2256,7 +2349,23 @@ async def diffusion_load(
# global checks. Refuse early so we do not first tear down an
# idle export checkpoint just to fail on the helper check inside
# load_model.
# Round 30 P1 #10: also publishes the public-load pending entry so
# a concurrent helper start cannot win the start lock between our
# snapshot and DiffusionBackend.load_model flipping is_loaded.
_raise_if_helper_advisor_busy("diffusion")
# Round 30 P1 #4: enforce the signed native_path_lease boundary the
# chat load path uses so local-path repo_id / base_repo cannot be
# probed without a frontend-issued grant. Hub ids pass through.
resolved_repo_id = _resolve_diffusion_repo_for_request(
payload.repo_id,
payload.native_path_lease,
operation = "load-diffusion-model",
) or payload.repo_id
resolved_base_repo = _resolve_diffusion_repo_for_request(
payload.base_repo,
payload.base_repo_native_path_lease,
operation = "load-diffusion-model",
)
# Round 18 P1 #3 + P1 #7: the route used to drop chat and idle
# export BEFORE ``backend.load_model`` ran its cheap validation
# (family inference, GGUF filename checks, gated-token failures,
@ -2274,9 +2383,9 @@ async def diffusion_load(
status = await asyncio.get_event_loop().run_in_executor(
None,
lambda: backend.load_model(
repo_id = payload.repo_id,
repo_id = resolved_repo_id,
gguf_filename = payload.gguf_filename,
base_repo = payload.base_repo,
base_repo = resolved_base_repo,
family_override = payload.family,
hf_token = payload.hf_token,
enable_model_cpu_offload = payload.enable_model_cpu_offload,
@ -2321,6 +2430,11 @@ async def diffusion_load(
except Exception as exc:
logger.exception("Diffusion load failed")
raise HTTPException(status_code = 500, detail = str(exc))
finally:
# Round 30 P1 #10: clear the public-load pending publish so a
# subsequent helper / advisor start can proceed once the
# diffusion load attempt has finished (success or failure).
_clear_public_load_window("diffusion")
@studio_router.post("/images/unload")

24
studio/backend/routes/training.py
View file

@ -127,6 +127,11 @@ async def start_training(
This endpoint initiates training in the background and returns immediately.
Use the /status endpoint to check training progress.
"""
# Round 30 P1 #7: track whether we published a public-load pending
# entry so the outer finally clears it on either success or
# failure (including any early HTTPException raised by the helper
# check itself).
training_load_window_published = False
try:
logger.info(f"Starting training job with model: {request.model_name}")
@ -272,6 +277,7 @@ async def start_training(
# the user's output artifact. Now we 409 first; the user
# stops the export and re-submits.
from routes.inference import (
_clear_public_load_window,
_raise_if_export_active,
_raise_if_helper_advisor_busy,
_release_chat_for,
@ -282,7 +288,13 @@ async def start_training(
_raise_if_export_active("training")
# Round 28 P1 #5: refuse before any release fires so AI Assist
# busy does not first tear down idle diffusion/export.
# Round 30 P1 #7: also publishes a public-load pending entry so
# a concurrent helper / advisor start cannot win the start
# lock between our snapshot and start_training flipping
# is_training_active. Paired clear lives in the outer
# ``finally`` below.
_raise_if_helper_advisor_busy("training")
training_load_window_published = True
# Round 18 P1 #8: release settled export FIRST so an export
# cleanup failure preserves the user's currently loaded chat
# model. The previous order (chat -> export) would drop chat
@ -336,6 +348,18 @@ async def start_training(
status_code = 500,
detail = f"Failed to start training: {str(e)}",
)
finally:
# Round 30 P1 #7: clear the public-load pending entry once the
# start attempt has finished. Skipped when the helper-busy
# check itself raised (no publish to clear) so the counter
# stays in sync with publishes.
if training_load_window_published:
try:
from routes.inference import _clear_public_load_window
except Exception:
pass
else:
_clear_public_load_window("training")
@router.post("/stop", response_model = TrainingStopResponse)

52
studio/backend/utils/datasets/llm_assist.py
View file

@ -46,11 +46,22 @@ README_MAX_CHARS = 1500
# * GPU : blocks public chat / training / export / diffusion loads
_HELPER_ADVISOR_CACHE_REFCOUNT: Counter[str] = Counter()
_HELPER_ADVISOR_GPU_REFCOUNT: Counter[str] = Counter()
# Round 30 P1 #7-#10: counter of public GPU workloads (chat /
# diffusion / training / export) that have passed the helper-busy
# snapshot but have not yet flipped their public ownership flags
# (``llama.is_loaded`` / ``loading_model_identifier`` /
# ``current_checkpoint`` / ``is_training_active``). Helper / advisor
# starts consult this so they cannot win the start lock and race a
# public load that already destroyed the previous owner.
_PUBLIC_LOAD_PENDING_COUNT: Counter[str] = Counter()
_HELPER_ADVISOR_LOCK = threading.Lock()
# Round 28 P1 #7 / #8 / #10: serialize helper / advisor STARTS so two
# concurrent invocations cannot both pass the busy precheck before
# either registers. Held only across the precheck + register window,
# not across the full helper run.
# Round 30 P1 #7-#10: public GPU loads also enter under this lock to
# publish their pending counter so a concurrent helper / advisor
# start sees the pending public owner and refuses VRAM.
_HELPER_ADVISOR_START_LOCK = threading.Lock()
@ -99,6 +110,38 @@ def _unregister_helper_advisor_repo(repo_id: str, *, gpu_owner: bool = True) ->
_HELPER_ADVISOR_GPU_REFCOUNT.pop(needle, None)
def _publish_public_load_pending(workload: str) -> None:
"""Mark a public GPU workload as mid-handoff. Must be called under
``_HELPER_ADVISOR_START_LOCK`` immediately after the helper-busy
snapshot succeeded (round 30 P1 #7-#10)."""
if not workload:
return
needle = workload.lower()
with _HELPER_ADVISOR_LOCK:
_PUBLIC_LOAD_PENDING_COUNT[needle] += 1
def _release_public_load_pending(workload: str) -> None:
"""Decrement the pending public-load counter once per matched
publish. Safe to call in finally even if the load failed."""
if not workload:
return
needle = workload.lower()
with _HELPER_ADVISOR_LOCK:
_PUBLIC_LOAD_PENDING_COUNT[needle] -= 1
if _PUBLIC_LOAD_PENDING_COUNT[needle] <= 0:
_PUBLIC_LOAD_PENDING_COUNT.pop(needle, None)
def public_load_pending() -> bool:
"""True if any public GPU workload has passed its helper-busy
snapshot but not yet flipped its public ownership flags. Helper /
advisor starts treat this as busy so they cannot race a public
load mid-handoff."""
with _HELPER_ADVISOR_LOCK:
return sum(_PUBLIC_LOAD_PENDING_COUNT.values()) > 0
def _strip_think_tags(text: str) -> str:
"""Strip <think>...</think> reasoning blocks emitted by some models.
@ -230,6 +273,15 @@ def _gpu_workload_busy_for_helper() -> bool:
"Skipping helper GGUF while another helper/advisor is using the GPU"
)
return True
# Round 30 P1 #7-#10: a public GPU load (chat / diffusion / training /
# export) that has passed its busy snapshot but not yet flipped its
# public ownership flags is still mid-handoff. Refuse so the helper
# does not race it for VRAM after the previous owner was torn down.
if public_load_pending():
logger.info(
"Skipping helper GGUF while a public GPU load is mid-handoff"
)
return True
if _diffusion_image_model_busy():
return True

36
studio/frontend/src/features/images/images-page.tsx
View file

@ -128,8 +128,11 @@ export function ImagesPage() {
const preset = CURATED_MODELS[presetIndex] ?? DEFAULT_PRESET;
const resolution = RESOLUTION_PRESETS[resolutionIdx];
const refreshStatus = useCallback(async () => {
setRefreshingStatus(true);
// Round 30 P2 #12: split the fetch from the spinner toggle so the
// mount + auto-poll effects can call the fetch without the
// synchronous setRefreshingStatus(true) that tripped
// react-hooks/set-state-in-effect.
const fetchAndUpdateStatus = useCallback(async () => {
try {
const next = await fetchDiffusionStatus();
setStatus(next);
@ -139,20 +142,25 @@ export function ImagesPage() {
lastErrorRef.current = msg;
toast.error("Could not fetch image-model status", { description: msg });
}
} finally {
setRefreshingStatus(false);
}
}, []);
const refreshStatus = useCallback(async () => {
setRefreshingStatus(true);
try {
await fetchAndUpdateStatus();
} finally {
setRefreshingStatus(false);
}
}, [fetchAndUpdateStatus]);
useEffect(() => {
// Round 30 P2 #12: defer the first refreshStatus call via
// queueMicrotask so the synchronous setRefreshingStatus(true)
// inside it does not trip the react-hooks/set-state-in-effect
// lint rule on the mount render.
queueMicrotask(() => {
void refreshStatus();
});
}, [refreshStatus]);
// Mount fetch goes through fetchAndUpdateStatus so the lint rule
// does not see any synchronous setState in the effect body; the
// user-driven Refresh button still calls refreshStatus to flip
// the spinner.
void fetchAndUpdateStatus();
}, [fetchAndUpdateStatus]);
// Round 27 P2: when the backend is mid-load (is_loading=true) the
// status label froze at "Loading..." until the user clicked
@ -161,10 +169,10 @@ export function ImagesPage() {
useEffect(() => {
if (!status?.is_loading) return;
const id = window.setInterval(() => {
void refreshStatus();
void fetchAndUpdateStatus();
}, 2000);
return () => window.clearInterval(id);
}, [status?.is_loading, refreshStatus]);
}, [status?.is_loading, fetchAndUpdateStatus]);
const handleLoad = useCallback(async () => {
setBusy("loading");