Studio: scale export GGUF size estimates from the real model size#6418
Conversation
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.
for more information, see https://pre-commit.ci
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a dynamic model-size estimation feature for the Export page, replacing hardcoded estimates with real, MoE-aware size calculations scaled from the model's FP16/BF16-equivalent size. On the backend, a new /api/models/export-size endpoint and caching mechanism were added, along with a suite of unit tests. On the frontend, a custom hook fetches these estimates to dynamically calculate and display human-readable sizes for each GGUF quantization level. The review feedback highlights three key improvements: changing the backend endpoint from async def to a synchronous def to prevent blocking the FastAPI event loop, updating the corresponding tests to call the endpoint synchronously, and clamping the index calculation in formatModelSize to a minimum of 0 to avoid out-of-bounds errors when formatting sizes less than 1 byte.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
|
|
||
|
|
||
| @router.get("/export-size", response_model = ExportSizeResponse) | ||
| async def get_export_size( |
There was a problem hiding this comment.
The get_export_size endpoint is defined as an async def but calls _export_size_cached, which performs blocking synchronous network and disk I/O (e.g., querying Hugging Face Hub via huggingface_hub and loading configurations via transformers). Because there are no await expressions inside this endpoint, defining it as async def will block the main asyncio event loop, preventing FastAPI from handling other concurrent requests.
Changing the endpoint definition to a standard synchronous def allows FastAPI to automatically run it in an external thread pool, keeping the event loop responsive.
| async def get_export_size( | |
| def get_export_size( |
| model = model, hf_token = None, current_subject = "test-user" | ||
| ) | ||
| ) | ||
|
|
||
| def test_known_model_returns_bytes_and_params(self): |
There was a problem hiding this comment.
Since get_export_size is being changed to a synchronous function to prevent blocking the event loop, calling it with asyncio.run in the test suite will raise a TypeError because it is no longer a coroutine. Update the test helper to call the endpoint synchronously.
return self.models_route.get_export_size(
model = model, hf_token = None, current_subject = "test-user"
)| const i = Math.min( | ||
| Math.floor(Math.log(bytes) / Math.log(1024)), | ||
| units.length - 1, | ||
| ); |
There was a problem hiding this comment.
If bytes is less than 1 (which can happen with fractional byte values or extremely small sizes), Math.floor(Math.log(bytes) / Math.log(1024)) will evaluate to a negative integer (e.g., -1). This results in an out-of-bounds index for the units array, returning a string like "512.0 undefined".
Clamping the index i to a minimum of 0 ensures that small byte values are formatted correctly as "B".
const i = Math.max(
0,
Math.min(
Math.floor(Math.log(bytes) / Math.log(1024)),
units.length - 1,
),
);![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3807fd1e7b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| signal?: AbortSignal, | ||
| ): Promise<ExportSizeEstimate> { | ||
| const response = await authFetch( | ||
| `/api/models/export-size?model=${encodeURIComponent(modelId)}`, |
There was a problem hiding this comment.
When the selected source is a private/gated Hugging Face model, this request only sends model, even though the backend endpoint accepts hf_token and the user can enter one on the Export page. In that scenario estimate_fp16_model_size_bytes cannot read safetensors/config metadata and returns null, so the new GGUF size estimates disappear for exactly the models that require authenticated access; include the token and make the hook refetch when it changes.
Useful? React with 👍 / 👎.
| resolved = model | ||
| if not is_local_path(model): | ||
| resolved = resolve_cached_repo_id_case(model) | ||
| fp16_bytes, total_params, source = _export_size_cached(resolved, hf_token) |
There was a problem hiding this comment.
When the selected model is not already resolved locally, _export_size_cached can perform blocking Hugging Face metadata/config requests via estimate_fp16_model_size_bytes; calling it directly from this async route ties up the server event loop for the duration of a slow Hub request, delaying unrelated API calls and live export/log endpoints. Run the estimate in a worker thread (for example with asyncio.to_thread) so a size hint cannot stall the rest of Studio.
Useful? React with 👍 / 👎.
| ) | ||
|
|
||
|
|
||
| @lru_cache(maxsize = 128) |
There was a problem hiding this comment.
Because the whole helper is memoized, transient failures are cached the same as successful estimates: if the first request for a model hits offline Hugging Face, a gated response before credentials are available, or any temporary exception, the (None, None, "unavailable") result is returned for subsequent selections of the same model/token until eviction or process restart. Cache only positive estimates (or use a short TTL for failures) so the Export page can recover once metadata becomes available.
Useful? React with 👍 / 👎.
- 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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b8757a1a71
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // The token lets the sizer read safetensors/config metadata for private or | ||
| // gated repos, which would otherwise return no estimate. | ||
| const tokenParam = hfToken | ||
| ? `&hf_token=${encodeURIComponent(hfToken)}` |
There was a problem hiding this comment.
When a user sizes a private or gated model, this puts their Hugging Face token directly into the request URL (hf_token=...). Query strings are routinely exposed in browser/devtool network logs and backend or proxy access logs, so this new size-estimation call can leak credentials; send the token in a request body or an internal header instead of appending it to the GET URL.
Useful? React with 👍 / 👎.
|
|
||
| Blocking (HF metadata / config / disk); call via a worker thread. | ||
| """ | ||
| key = (model, hf_token) |
There was a problem hiding this comment.
When a private or gated model estimate succeeds, this module-global cache stores the full Hugging Face token as part of the key for the lifetime of the backend process. The cache only needs to separate credential scopes, so retaining raw secrets in long-lived memory unnecessarily exposes them in debug dumps or memory snapshots; key by a token fingerprint or another non-secret discriminator instead.
Useful? React with 👍 / 👎.
| ``estimate_fp16_model_size_bytes`` (safetensors -> config -> local -> vllm). | ||
| """ | ||
| resolved = model | ||
| if not is_local_path(model): |
There was a problem hiding this comment.
When an authenticated caller passes an existing broad local directory such as / or a home directory, this branch leaves it as the model to size, and _export_size_cached ultimately reaches the local weight sizer that recursively walks the directory looking for weight files. A single malformed export-size request can therefore burn I/O over the whole filesystem and tie up worker threads; restrict this endpoint to known model directories or reject overly broad local paths before invoking the sizer.
Useful? React with 👍 / 👎.
- 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.
for more information, see https://pre-commit.ci
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ecf1699a9f
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| root_resolved = root.resolve() | ||
| except (OSError, RuntimeError, ValueError): | ||
| continue | ||
| if target == root_resolved or target.is_relative_to(root_resolved): |
There was a problem hiding this comment.
When model is exactly studio_root(), outputs_root(), exports_root(), or cache_root(), this equality branch returns true; _export_size_cached then calls the local weight sizer, which recursively walks the whole tree. A malformed authenticated /api/models/export-size?model=$STUDIO_HOME or cache-root request can still burn I/O over all Studio outputs/caches instead of one model directory. Fresh evidence: the new allowlist still admits the broad root itself via target == root_resolved; require a specific model directory before sizing.
Useful? React with 👍 / 👎.
| return False | ||
| if not target.exists(): | ||
| return False | ||
| for root in (studio_root(), outputs_root(), exports_root(), cache_root()): |
There was a problem hiding this comment.
When a user directly exports a local model that Studio already discovered from LM Studio or a registered custom scan folder, the frontend sends that local path to /export-size, but this allowlist only accepts Studio roots. Those selections can still export, yet the endpoint returns unavailable, so the new quant-size labels and summary disappear for a supported local-model flow; include the same trusted LM Studio/custom roots used by the local model picker, while still rejecting arbitrary paths.
Useful? React with 👍 / 👎.
| cached = _EXPORT_SIZE_CACHE.get(model) | ||
| if cached is not None: | ||
| return cached |
There was a problem hiding this comment.
When a private or gated model is first sized with X-HF-Token, the successful result is cached only by model, so a later request for the same repo without that token (or from another authenticated subject) returns the privileged cached size instead of the unavailable result the sizer would produce. That leaks private model metadata and makes the UI show estimates under credentials that cannot access the repo; key successful remote estimates by a non-secret credential scope such as subject plus token fingerprint, or only reuse unauthenticated successes globally.
Useful? React with 👍 / 👎.
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
- 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.
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).
danielhanchen
force-pushed
the
export-size-fix
branch
from
ada94f6 to
56d9824
Compare
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
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.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: d1ceaa82d2
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| // Real (MoE-aware) fp16 size, used to scale the GGUF quant estimates. | ||
| const { fp16Bytes } = useExportSizeEstimate( | ||
| sourceBaseModelName, |
There was a problem hiding this comment.
When exporting a checkpoint from a full fine-tune whose base_model is a local/custom path outside the Studio roots (or has since been moved), this hook asks /export-size to size sourceBaseModelName rather than the selected checkpoint directory. That makes the endpoint return unavailable even though the checkpoint itself is under outputs_root() and could be sized locally, so GGUF quant labels and the summary disappear for a supported checkpoint export flow; use the selected checkpoint path for non-adapter checkpoint exports, keeping the base-model estimate for LoRA adapters.
Useful? React with 👍 / 👎.
| try: | ||
| from utils.hardware.hardware import estimate_fp16_model_size_bytes | ||
|
|
||
| fp16_bytes, source = estimate_fp16_model_size_bytes(model, hf_token = hf_token) |
There was a problem hiding this comment.
When model is a local LoRA adapter inside an allowed Studio root, this call can still end up sizing an arbitrary path because estimate_fp16_model_size_bytes first resolves adapters through base_model_name_or_path and then the local weight sizer recursively walks that resolved base. A malformed adapter under outputs/cache with a base like / or a home directory bypasses _is_sizable_local_path and can burn I/O over the filesystem; re-check the resolved identifier before allowing local sizing.
Useful? React with 👍 / 👎.
| if candidate == root or candidate.startswith(root + os.sep): | ||
| # Contained in a trusted root; safe to touch the filesystem. | ||
| return os.path.exists(candidate) |
There was a problem hiding this comment.
When an allowed Studio path is itself a directory symlink to somewhere outside the Studio roots (for example outputs/link -> /), this lexical prefix check still accepts it and the subsequent local weight sizer follows that symlink while recursively walking the candidate. That bypasses the local-path allowlist and lets a malformed export-size request scan arbitrary filesystem trees; compare resolved real paths or reject symlinked directory candidates before sizing.
Useful? React with 👍 / 👎.
| try: | ||
| from utils.hardware.hardware import estimate_fp16_model_size_bytes | ||
|
|
||
| fp16_bytes, source = estimate_fp16_model_size_bytes(model, hf_token = hf_token) |
There was a problem hiding this comment.
When a user directly sizes a Studio training-run directory that contains checkpoint-* subdirectories, this passes the run root to the estimator; its local fallback recursively sums every weight file below that root and then prefers the larger local byte count over the config-derived size. The GGUF estimate can therefore be multiplied by the number of saved intermediate checkpoints even though export loads only the model at the selected root, so local full-finetune direct exports show inflated quant sizes; size only the selected model files/snapshot instead of the whole tree.
Useful? React with 👍 / 👎.
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.
|
You have reached your Codex usage limits for code reviews. You can see your limits in the Codex usage dashboard. |
for more information, see https://pre-commit.ci
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
| except (OSError, RuntimeError, ValueError): | ||
| continue | ||
| if real == real_root or real.startswith(real_root + os.sep): | ||
| return os.path.exists(real) |
3 participants
Summary
The Export page showed hardcoded, model-independent GGUF quant sizes (Q8_0 ~8.2 GB, BF16 ~14.2 GB, ...) calibrated for an ~8B model. For a 35B MoE model such as Qwen3.6-35B-A3B (67 GiB bf16, Q8 ~34 GiB) the picker reported Q8 ~8.2 GB, which is wrong and confusing. Only the displayed estimate was affected; the actual export via
save_pretrained_ggufwas always correct.This scales the estimate from the model's real, MoE-aware fp16 size.
Changes
Backend
GET /api/models/export-size, returning a model's fp16/bf16-equivalent size and total params, reusing the existing MoE-awareestimate_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.ExportSizeResponseschema andtests/test_export_size_estimate.py.Frontend
QUANT_OPTIONSsizes with a bits-per-weight map; each quant scales from the real fp16 size (bytes ~= fp16_bytes * bpw / 16), formatted in GiB to match the model selector.useExportSizeEstimatehook fetches the size when a checkpoint or model is selected; sizes are hidden when unknown rather than shown as a fixed number.Verification
unsloth/Qwen3.6-35B-A3B: Q8 ~36 GiB (was ~8.2 GB), within ~3.5% of the published GGUF; Q4_K_M within ~2%.Before/after for Qwen3.6-35B-A3B: