From 79da5d910dffac61158049b884a87a54959f0b9b Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 13:43:51 +0000 Subject: [PATCH] Fix/adjust diffusion: round 27 follow-up P1 batch for PR #5754 Five additional P1 findings round 27 reviewer flagged on top of the round 27 commit 6c528fb0 (Counter refcount + handoff visibility were already covered). Three remaining studio.txt / no-torch-runtime hub suggestions are NOT applied because they would re-break CI; the empirical evidence (round 26 commit 65ea3a2c restored CI green) takes precedence over the reviewer's stale-state suggestion. 1. models/training.py TrainingStartRequest: extend the embedded HF token validator to subset, train_split, eval_split. Round 26 only added the control-char guard to those three; the token guard was asymmetric and would accept owner/data\\nFAKE hf_abcdef... payloads through subset / split fields. 2. models/datasets.py CheckFormatRequest: extend both validators (control chars + embedded HF token) to subset and train_split. Same asymmetric-fix bug as #1. 3. models/data_recipe.py SeedInspectRequest: extend both validators to subset and split. Same pattern. 4. utils/datasets/llm_assist.py precache_helper_gguf: register the helper repo in the helper/advisor refcount registry around the hf_hub_download loop, then unregister in the finally. Without this, the FastAPI-startup background pre-cache could be racing a concurrent DELETE /api/models/delete-cached against the same cache directory. The runtime helper / advisor calls already register (round 26 P1 #13/#14) but the precache was the asymmetric gap. 5. routes/models.py _loaded_model_matches_deleted_path: match bidirectionally (active under target OR target under active) so deleting a child directory of a loaded local model (.../my-flux/ text_encoder while .../my-flux is loaded) trips the guard. Mirrors the diffusion delete-guard symmetric path-overlap check. Tests: 105 targeted (diffusion + cache + inference_validation) and the broader backend suite pass locally. --- studio/backend/models/data_recipe.py | 6 ++++-- studio/backend/models/datasets.py | 6 ++++-- studio/backend/models/training.py | 6 +++++- studio/backend/routes/models.py | 16 +++++++++++++--- studio/backend/utils/datasets/llm_assist.py | 6 ++++++ 5 files changed, 32 insertions(+), 8 deletions(-) diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index 06fdcb0963..2e538e138c 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -92,12 +92,14 @@ class SeedInspectRequest(BaseModel): # Round 26 P1 #11: dataset_name reaches HF + log/echo paths, so # mirror the hardening other dataset request models already do. - @field_validator("dataset_name") + # Round 27 P1 #7: split and subset also flow into HF dataset + # APIs / errors and must be guarded the same way. + @field_validator("dataset_name", "subset", "split") @classmethod def _no_dataset_name_control_chars(cls, v, info): return _no_control_chars(v, info.field_name) - @field_validator("dataset_name") + @field_validator("dataset_name", "subset", "split") @classmethod def _no_dataset_name_embedded_hf_tokens(cls, v, info): return _reject_embedded_hf_token(v, info.field_name) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 6f9de26939..28a4016514 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -32,12 +32,14 @@ class CheckFormatRequest(BaseModel): values.setdefault("train_split", values.pop("split")) return values - @field_validator("dataset_name") + # Round 27 P1 #6: subset / train_split also flow into HF dataset + # APIs and errors/responses, so they need the same hardening. + @field_validator("dataset_name", "subset", "train_split") @classmethod def _no_dataset_name_control_chars(cls, v, info): return _no_control_chars(v, info.field_name) - @field_validator("dataset_name") + @field_validator("dataset_name", "subset", "train_split") @classmethod def _no_dataset_name_embedded_hf_tokens(cls, v, info): return _reject_embedded_hf_token(v, info.field_name) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 234474fd23..de440a6c06 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -67,7 +67,11 @@ class TrainingStartRequest(BaseModel): def _no_model_name_control_chars(cls, v, info): return _no_control_chars(v, info.field_name) - @field_validator("model_name", "hf_dataset") + # Round 27 P1 #2: subset / train_split / eval_split are reflected + # in status + error messages and persisted to job records, so the + # embedded-token guard must cover them too. Round 26 only added + # the control-char guard to those three. + @field_validator("model_name", "hf_dataset", "subset", "train_split", "eval_split") @classmethod def _no_model_name_embedded_hf_tokens(cls, v, info): return _reject_embedded_hf_token(v, info.field_name) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 8b3025e056..3ec3a9372e 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1795,7 +1795,15 @@ def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) -> try: active = Path(active_model).expanduser().resolve() target = deleted_path.resolve() - return active == target or (target.is_dir() and active.is_relative_to(target)) + # Round 27 P1 #8: match bidirectionally so deleting a child + # directory of a loaded local model (e.g. .../my-flux/text_encoder + # while .../my-flux is loaded) also trips the guard. Mirrors + # the diffusion delete-guard pattern. + return ( + active == target + or (target.is_dir() and active.is_relative_to(target)) + or (active.is_dir() and target.is_relative_to(active)) + ) except (OSError, RuntimeError, ValueError) as e: logger.debug( "Could not resolve loaded/deleted model paths; falling back to string comparison: %s", @@ -1803,8 +1811,10 @@ def _loaded_model_matches_deleted_path(active_model: str, deleted_path: Path) -> ) active_lower = active_model.lower() target_lower = str(deleted_path).lower() - return active_lower == target_lower or active_lower.startswith( - f"{target_lower}{os.sep}" + return ( + active_lower == target_lower + or active_lower.startswith(f"{target_lower}{os.sep}") + or target_lower.startswith(f"{active_lower}{os.sep}") ) diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 37c17329c8..d6dea74c47 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -127,6 +127,11 @@ def precache_helper_gguf(): "UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT ) + # Round 27 P1 #4: register the repo so DELETE /api/models/delete-cached + # cannot rmtree the cache directory while we are mid-download. Helper + # / advisor runtime calls already register, but the startup precache + # was the asymmetric gap that let cache delete race the first download. + _register_helper_advisor_repo(repo) try: from huggingface_hub import HfApi, hf_hub_download from huggingface_hub.utils import disable_progress_bars, enable_progress_bars @@ -158,6 +163,7 @@ def precache_helper_gguf(): except Exception as e: logger.warning(f"Failed to pre-cache helper GGUF: {e}") finally: + _unregister_helper_advisor_repo(repo) try: enable_progress_bars() except Exception as e: