unsloth/studio/backend/utils/hardware
Daniel Han 22e6d64493
Studio: scale export GGUF size estimates from the real model size (#6418)
* Studio: scale export GGUF size estimates from the real model size

The Export page showed hardcoded, model-independent GGUF quant size
labels (Q8_0 ~8.2 GB, BF16 ~14.2 GB, ...) calibrated for an ~8B model.
For a 35B MoE model like Qwen3.6-35B-A3B (67 GiB bf16, Q8 ~34 GiB) the
picker wrongly reported Q8 ~8.2 GB. Only the displayed estimate was
wrong; the actual export via save_pretrained_gguf was always correct.

Add GET /api/models/export-size, which returns a model's MoE-aware
fp16/bf16-equivalent size and total params using the existing
estimate_fp16_model_size_bytes (safetensors -> config -> local -> vllm).
The result is memoized and degrades to nulls so a size hint can never
break the Export page.

The Export picker now scales each quant from that size
(bytes ~= fp16_bytes * bits_per_weight / 16, GiB units to match the
model selector), and renders no size when it is unknown rather than a
misleading fixed number. The Est. size summary in the page and dialog
is restored now that the value comes from the backend.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio export-size: address review feedback

- Run the size estimate off the event loop with asyncio.to_thread so a slow
  Hugging Face request cannot stall other API or SSE endpoints.
- Cache only successful estimates; a transient failure (offline, gated before
  credentials) is no longer pinned as unavailable until restart.
- Forward the HF token so private and gated models can be sized, and refetch
  when the token changes.
- Clamp the size formatter index so sub-1-byte values cannot pick an
  out-of-range unit.

* Studio export-size: address second review pass

- Send the HF token in an X-HF-Token header instead of the query string, so
  it never lands in URLs, logs, or browser history.
- Key the estimate cache by model id only (the fp16 size is token independent),
  so HF tokens are never retained in the cache.
- Restrict local-path sizing to known Studio roots (outputs/exports/cache/home)
  so an authenticated caller cannot trigger a scan of an arbitrary directory.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio export-size: fix CI (import-hoist + isolated-load test stubs)

- Import ExportSizeResponse from models.models in routes/models.py instead of
  re-exporting it through models/__init__.py, so the import-hoist lint does not
  flag a newly added but un-loaded re-export (models/__init__.py is unchanged).
- Add Header and ExportSizeResponse to the stubbed fastapi / models.models in
  test_export_absolute_paths.py, which loads routes/models.py in isolation.

* Studio: validate export-size local path before filesystem access

CodeQL flagged the export-size local-path guard as path injection: the
user-provided model path was resolved and stat-ed before it was checked
for containment under a Studio data root. Decide containment by lexical
normalization (normpath/abspath/expanduser, no filesystem access) and
only touch the filesystem once the path is proven to sit under a trusted
root, so an unvalidated value never reaches a filesystem call. Add a
direct containment unit test (under-root, root itself, missing, /etc,
and '..' traversal).

* Studio: trim export-size comments to be more concise

Shorten docstrings and comments on the export-size endpoint, helpers, tests,
and frontend size utilities; drop comments that just restate the code. Verified
code-identical (comments only) via AST/TS-compiler check. No behavior change.

* Studio: harden export-size local-path handling

Address review feedback on the export-size endpoint's local sizing:
- Resolve symlinks and re-verify containment in _is_sizable_local_path so a
  symlink inside a Studio root can't point the sizer outside it.
- Re-validate the resolved LoRA base before sizing, so a crafted adapter
  whose base_model points outside the roots can't redirect the scan.
- Skip nested checkpoint-*/global_step* snapshots when summing local weight
  sizes so a run dir's intermediate checkpoints don't inflate the estimate.
- Size the checkpoint directory for full fine-tune checkpoint exports (whose
  base may be a local/custom path), keeping base-model sizing for adapters.

Adds tests for the adapter-base escape, symlink escape, and nested-checkpoint
exclusion.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
2026-06-18 05:44:17 -07:00
..
__init__.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
amd.py Windows/WSL installer: fix winget msstore cert failure, amd-smi DiskPart prompt, and enable AMD GPU (Strix Halo gfx1151) (#5940) 2026-06-10 04:24:49 -07:00
apple.py Guard Apple GPU power against negative counter-reset readings (#6235) 2026-06-12 01:56:05 -07:00
hardware.py Studio: scale export GGUF size estimates from the real model size (#6418) 2026-06-18 05:44:17 -07:00
nvidia.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00
VRAM_ESTIMATION.md Update VRAM estimator to cater to broader model configs (#5175) 2026-05-05 04:12:36 -07:00
vram_estimation.py Reduce and tighten code comments and docstrings repo-wide (#6095) 2026-06-08 23:09:51 -07:00