Commit graph

6 commits

Author SHA1 Message Date
Daniel Han
187144d4e7
Reduce and tighten code comments and docstrings repo-wide (#6095)
Trim and tighten code comments and docstrings across the repository. Comment-only: every changed file verified code-identical to main via AST/token comparison.
2026-06-08 23:09:51 -07:00
Daniel Han
8292e699e4
Studio: make code comments and docstrings more succinct (#6029)
Trim and tighten code comments and docstrings across studio/ Python. Comment-only: every changed file verified code-identical to main via AST/token comparison.
2026-06-08 23:07:28 -07:00
Daniel Han
3ce187da02
Formatting: ruff line-length 100, kwarg-spacing passes, drop blank after short local imports (#6079)
Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
2026-06-08 04:24:13 -07:00
Wasim Yousef Said
e35cbfb454
Add native GGUF intake to Studio (#5246)
* feat(studio): add Tauri native GGUF intake

* feat(studio): polish native GGUF intake

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

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

* fix(studio): load backend helpers during local setup

* fix(studio): acquire native load lease before unload

* Studio: harden native path lease verification and Tauri intake

- Wrap path.resolve(strict=True) and Path.stat() in NativePathLeaseError so a deleted or unmounted GGUF returns 400 instead of leaking the full filesystem path through the generic load_model/validate_model handler.
- Re-apply _reject_network_or_device_path to the resolved canonical path for defense in depth after symlink resolution.
- Replace try/except ValueError pattern in the device-path guard with Path.is_relative_to; the previous shape silently swallowed NativePathLeaseError (which subclasses ValueError) so /dev,/proc,/sys were never actually rejected.
- Broaden the lease redaction regex and dict-key check (Python and Rust diagnostics) to cover both native_path_lease and nativePathLease so the camelCase form emitted by Tauri/frontend payloads is also redacted.
- Hoist the redact_native_paths import to module top in loggers/handlers; the recursive filter no longer pays a per-record import lookup.
- Persist activeNativePathToken in the chat runtime store so the rollback branch can mint a fresh lease and reload the previous native GGUF when a new load fails after unload; clear it in clearCheckpoint and overwrite it on each successful load.
- use-native-drop: read options through a ref so the Tauri onDragDropEvent listener is registered once and stays attached across option changes; reject ambiguous multi-file drops up front instead of silently registering only the first GGUF.
- pick_native_model: use an async pick_file with a tokio oneshot channel instead of blocking_pick_file so the Tokio worker is not held for the duration of the OS dialog.
- registerNativeModelPath: drop the duplicate sourceKind argument; the Rust command parameter is source_kind.
- install_python_stack: insert the script directory (studio/) on sys.path; the previous insert pointed at studio/backend/ which does not satisfy `from backend.utils.wheel_utils import ...`.

* install_python_stack: keep _BACKEND_DIR on sys.path

Restore the studio/backend insertion. Although the immediately following `from backend.utils.wheel_utils import (...)` is satisfied by studio/ already being on sys.path[0] when invoked as `python studio/install_python_stack.py`, wheel_utils itself runs `from utils.native_path_leases import ...`, which requires studio/backend/ to be importable. Without the backend insertion, the existing tests/python/test_install_python_stack.py collection fails with ModuleNotFoundError: No module named 'utils'.

* Studio: tighten native path lease lifecycle and Tauri intake IPC

- register_native_model_path now hardcodes NativePathSourceKind::Drop on the Rust side and the frontend stops sending source_kind. The previous JS payload (source_kind only) never reached the Rust deserializer because Tauri's default ArgumentCase::Camel maps the Rust parameter source_kind to the JS key sourceKind, so drag/drop registration silently failed. Hardcoding the source kind also keeps audit metadata trustworthy on this command.
- Add native_path_secret_removed_for_child_start context manager and wrap multiprocessing.Process.start() at the inference, export, training, and data-recipe job spawn sites. The previous wrapper-only scrub left UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET visible to spawn-platform import-time worker code. The wrapper run_without_native_path_secret stays as defense-in-depth inside the child.
- Stop passing exc_info=True from the native-grant load/validate error logs in routes/inference.py. The structlog filter_sensitive_data processor runs before the renderer, so ConsoleRenderer formatted tracebacks bypassed redaction; the redacted str(e) preserves the message text.
- Replace the os.path.normcase string equality on the resolved canonical path with Path.samefile (with a normcase fallback) so Windows leases that differ only in extended-length \\?\ prefix or short-name spelling are accepted.
- Wrap consumeNativePathToken in its own try/catch in the chat runtime rollback. If the previous native-model token has aged out of TOKEN_TTL we now surface a clear modelsError instead of silently swallowing the rollback inside the outer catch.
- Reject non-ASCII lease strings in _split_lease and convert UnicodeEncodeError / binascii.Error / ValueError raised by _b64decode into NativePathLeaseError so verify_native_path_lease never escapes raw exceptions to the route handler.
- Tighten dropStateForPaths to mark multi-file payloads invalid so the overlay matches the post-fix drop handler that rejects the same payload.
- Replace the one-shot fetch in useNativePathLeasesSupported with a delayed-retry loop so the picker/drop becomes available once the backend is up rather than staying disabled for the rest of the session after a transient failure.
- Drop the unused setActiveNativePathToken setter; the value is set via setState directly in use-chat-model-runtime.
- Add a toast on auto-load failure in use-native-drop so a collapsed model selector does not hide the error.
- Burn the lease nonce before _validate_current_stat so a stat-failed lease is single-use even if a later state change happens to match the original size/mtime.

* Studio: cache lease secret, harden native path stat checks, polish intake UX

- Cache the decoded UNSLOTH_STUDIO_NATIVE_PATH_LEASE_SECRET on first verify and validate that it is base64-decodable and at least 32 bytes. Subsequent _decode_secret calls return from the cache and never touch os.environ, so concurrent /api/inference/load and /api/health requests no longer race with native_path_secret_removed_for_child_start scrubbing the env. native_path_leases_supported now wraps _decode_secret so the health flag matches what verify_native_path_lease actually accepts.
- Replace path.is_file()/is_dir() + path.stat() with os.lstat() in _validate_current_stat and explicitly reject S_ISLNK; size and mtime checks now refer to the link itself, closing the same-size+same-mtime symlink-swap window that the prior follow-symlink stat() left open.
- Add an issued_at_ms < expires_at_ms sanity check in _validate_payload to reject internally inconsistent (HMAC-protected) lease payloads.
- Sort _NATIVE_PATH_REDACTIONS by length (descending) before iterating in redact_native_paths so a longer registered path is replaced before a shorter prefix path; otherwise logs containing /foo/X.gguf.bak after only /foo/X.gguf was registered would leak the .bak suffix.
- classify_existing_path now re-checks the canonical path with symlink_metadata after canonicalize, so a regular file that is replaced with a symlink in the small canonicalize window is rejected at registration.
- ModelSelector renders the local file picker as its own block (not in the eject ternary), so a user with an active model can still replace it via the picker rather than only via drag/drop.
- useNativePathLeasesSupported caps the readiness probe at MAX_READINESS_POLLS (60 = ~5 minutes) and aborts the in-flight fetch on unmount via AbortController, so a permanently-disabled backend stops generating sustained traffic and hot-reload no longer leaks open connections.
- useChooseNativeModel returns a stable useCallback closure and guards the OS dialog with a useRef so rapid double-clicks cannot open multiple dialogs and orphan Rust tokens.
- Branch the multi-file drop toast: if no GGUF was present we say "Only .gguf model files can be dropped here." and otherwise "Drop a single .gguf model file." so users dropping non-GGUF attachments get an accurate explanation.

* native_path_leases: lstat the signed canonical path before resolving

The earlier change to lstat inside _validate_current_stat operates on grant.canonical_path, which is the post-resolve target. If the user atomically replaces the originally-signed file with a symlink to a different file of identical size and mtime, path.resolve(strict=True) follows the symlink, samefile returns True (both ends share the new inode), and the lstat in _validate_current_stat sees the regular target file rather than the symlink, so the swap goes undetected.

Add an os.lstat on the signed canonical path before path.resolve(strict=True), and reject S_ISLNK there. The lstat in _validate_current_stat stays as defense-in-depth for swaps that occur strictly between resolve and stat.

* Studio: scrub native lease secret before mp.Queue spawn and tighten lease lifecycle

- Move _CTX.Queue / _CTX.Event / _CTX.Process construction inside native_path_secret_removed_for_child_start at the inference, export, training and data-recipe spawn sites. The first Queue creation lazily spawns Python's multiprocessing.resource_tracker child, so when it ran outside the scrub context the tracker process inherited the lease secret. Reproduced via the proc filesystem environ entry; the wrapped order keeps the tracker clean.
- native_path_secret_removed_for_child_start now refcounts entries: the env var is popped on the first entry and restored only when the last context exits. Concurrent training/inference/export starts no longer serialize on the env lock across the entire proc.start yield, while still guaranteeing the env stays empty for the duration of every overlapping spawn.
- run_without_native_path_secret now also nulls the module-level cached lease secret. With the existing spawn-only multiprocessing context the cache is irrelevant in practice, but a future fork caller would otherwise inherit the in-memory secret even though the env var was scrubbed.
- filter_sensitive_data now applies the native lease key check on the top-level event_dict, not only on nested dicts, so a logger call that includes a lease value as a top-level keyword field actually redacts it (the bare value does not match the prefix-anchored regex).
- chat-page loadNativeModelIntent now passes intent.id to clearModelIntent so a second drag-drop during an in-flight first auto-load is not wiped from the chip area when the first resolves.
- Bump useNativePathLeasesSupported's MAX_READINESS_POLLS from 60 to 720 so first-run installs that compile llama.cpp from source or download large CUDA wheels (well past 5 minutes) don't permanently disable the native picker.

* native_path_leases: serialize first-decode against scrub context

_decode_secret used a separate _SECRET_INIT_LOCK from the env scrub's _NATIVE_PATH_ENV_LOCK, so the very first decode (before the cache is populated) could race a concurrent native_path_secret_removed_for_child_start and read os.environ during the env-empty window, raising "Native path grants require the managed desktop backend." Subsequent calls hit the cache and were already safe.

Acquire _NATIVE_PATH_ENV_LOCK around the env read inside _SECRET_INIT_LOCK and fall back to _SCRUB_SAVED_SECRET when the scrub has temporarily popped the env var. Lock ordering (init then env) is consistent with no other caller, so no deadlock.

* Studio: surface native model load errors and harden native path label cache

- Native model load and validate now bubble up the actual exception (with
  paths redacted) and apply the same friendly-error rewrite the non-native
  path uses, so users see "CUDA OOM", "trust_remote_code required", etc.
  instead of a generic "Failed to load native model: <label>".
- run_without_native_path_secret now also nulls _SCRUB_SAVED_SECRET so a
  forked grandchild that imports native_path_leases cannot recover the
  secret via the scrub-aware fallback in _decode_secret.
- _NATIVE_PATH_LABELS now has its own 10000-entry cap independent of the
  100-entry redaction list, so display_label_for_native_path no longer
  falls back to returning the raw canonical path after 101 native paths
  in one session. Redaction list keeps the 100-entry cap for log-scan
  performance.
- _validate_payload now also rejects null bytes in display_label, which
  is echoed back in HTTP responses and log lines.

* Studio: harden native path lease validation and chained native rollback

- child_env_without_native_path_secret now copies os.environ under
  _NATIVE_PATH_ENV_LOCK so a concurrent scrub-context env pop cannot
  raise RuntimeError: dictionary changed size during iteration in a
  background hardware scan or other env reader.
- _validate_payload and grant construction route every signed numeric
  field (version, issued_at_ms, expires_at_ms, size_bytes, modified_ms)
  through new _required_int / _optional_int helpers that wrap raw int()
  ValueError into NativePathLeaseError. The single upstream catcher
  produces 400 instead of 500 for malformed signed payloads.
- verify_native_path_lease now runs _validate_current_stat before
  _consume_nonce, so a transient stat error on the canonical path no
  longer permanently burns the nonce. Concurrent verifies still
  serialize through _consume_nonce, so single-use is preserved.
- Chained native model rollback now restores activeNativePathToken in
  the chat runtime store after a successful rollback loadModel. Without
  this, a second consecutive failed switch could not re-roll-back
  because the store token had been overwritten by the failed attempt.
- validate_model now applies the same not_supported_hints friendly
  rewrite to native model errors that load_model already does, so a
  native .gguf that fails validation with an upstream "is not supported"
  message gets the same actionable wording as the non-native branch.

* Studio: harden native path log redaction, status disclosure, and chip lifecycle

- structlog processor chain now runs format_exc_info before
  filter_sensitive_data so traceback strings are produced (and then
  redacted) rather than passed through as untouched (type, value, tb)
  tuples that the JSON or console renderer formats after the redaction
  filter has already finished.
- native_path_secret_removed_for_child_start clears _CACHED_LEASE_SECRET
  in addition to popping the env var, so a fork during the scrub window
  cannot inherit the cached bytes via the parent's heap. Parent verify
  calls during the window keep working through the existing scrub-aware
  fallback in _decode_secret.
- load_model's except ValueError handler now redacts native paths and
  uses the native model log label when native_grant_backed is true.
  Previously a ValueError raised after lease verification (e.g. from
  ModelConfig.from_identifier or downstream GGUF parsing) returned the
  raw exception string in the HTTP response body.
- llama_cpp_backend now records the native display label at GGUF load
  time, and /api/inference/status prefers it over the redaction store.
  After a Python backend restart the redaction store is empty; the
  attribute keeps the friendly label, and an absolute model_identifier
  with no other label source falls back to the basename so the canonical
  path no longer appears in active_model.
- reveal_path_token uses native "reveal and select" commands on macOS
  (open -R) and Windows (explorer /select,) so the file is highlighted
  in the file manager. Linux keeps the existing parent-directory open.
- Native model rollback that fails because the previous token cannot be
  consumed now throws a rollback-specific Error, and the outer empty
  catch was replaced with one that re-throws the rollback error. The
  rollback-specific message now reaches the user instead of being
  overwritten by the original load error message.
- NativeModelChip tracks the Rust token's expiresAtMs on a single
  setTimeout, disables the Load button at expiry, and relabels it
  "Select again" with an explanatory tooltip so users do not click into
  a guaranteed-failure path after the 15-minute TTL elapses.

* Studio: tighten native artifact policy, mmproj sibling check, and intake UX

- is_open_safe_artifact no longer grants Open for directories. Reveal
  already handles directory navigation, so the change closes the
  attack surface where a macOS .app artifact could be launched via
  open_path_token + open::that_detached.
- Display labels are sanitized in classify_existing_path. Control
  characters in filenames (newlines, tabs, NUL et al.) are replaced
  with spaces and the label is trimmed and capped, so a file named
  with embedded newlines cannot inject forged log lines or scramble
  the UI status panel.
- validate_entry_path skips the size_bytes/modified_ms equality check
  when the operation is Reveal or Open. Cloud-sync agents (Dropbox,
  iCloud Drive, OneDrive) routinely rewrite extended-attribute
  metadata which bumps mtime, and the user expects Reveal/Open to
  remain available for files in synced folders.
- llama_cpp_backend gains a _native_grant_backed flag at GGUF load
  success. /api/inference/status only applies the absolute-path
  basename fallback when that flag is true, so a non-native absolute
  local GGUF still reports its canonical model_identifier and unload
  by identifier keeps working.
- Native vision GGUFs now run through _validate_native_mmproj_companion
  before llama-server starts: the companion mmproj must be a regular
  file, not a symlink, and must live in the same resolved directory as
  the granted GGUF. This stops a hostile sibling or symlinked mmproj
  from being loaded under a single-file lease.
- Chained native rollback restructured: the rollback loadModel + state
  + refresh runs inside its own try/catch that swallows so the outer
  throw error surfaces the ORIGINAL load failure. The native-token
  consume-failure case still throws the rollback-specific message
  early, before the inner block runs, so its actionable guidance is
  preserved.
- Loading-model state and the duplicate-load guard in the chat runtime
  hook now compare both the model id and the native path token. Two
  drops or picks with the same basename in different folders no longer
  silently dedup; the second token is honored.
- chat-page loadNativeModelIntent awaits selectModel before clearing
  the pending intent. If selectModel returns early via dedup or
  throws, the chip and its token stay so the user can retry instead
  of losing the selection.
- NativeModelChip's Reveal button is disabled when the lease has
  expired (Rust would reject it anyway), and the Load button label
  reads "Expired" instead of "Select again" so the disabled element
  no longer promises an action it cannot perform.

* [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: Daniel Han <danielhanchen@gmail.com>
2026-05-04 11:46:18 +02:00
Wasim Yousef Said
a5eb2e3d50
Add tauri (#5144)
* add unsloth studio desktop app

* Fix review findings

- studio/src-tauri/tauri.conf.json: retarget updater to staging repo
  (danielhanchen/unsloth-staging-2); switch to unslothai/unsloth on upstream merge.
- studio/src-tauri/linux/postremove.sh: drop the interactive read loop and the
  /home/* iteration. Package maintainer scripts must stay non-interactive and
  must not touch other users' data.
- studio/frontend/src/app/auth-guards.ts: honor tauriAutoAuth() boolean. Failed
  auto-auth now redirects to /login; requireGuest/requirePasswordChangeFlow
  only redirect to /chat when auth succeeds. The new early-return on failed
  auth is intentional so the login / change-password flows remain reachable
  when desktop auth is not yet established.
- studio/frontend/src/config/env.ts: keep fetched=false on health failure so
  later calls retry instead of caching the client-side platform guess.
- studio/src-tauri/src/install.rs: pick the available system package manager
  (apt-get, dnf, zypper, pacman); AppImage bundles run on non-Debian distros.
- studio/frontend/src/lib/open-link.ts + markdown-text/sources callers: return
  boolean from openLink so callers only preventDefault on handled URLs; relative
  hrefs now navigate natively.
- studio/frontend/src/features/settings/tabs/about-tab.tsx: fetch(apiUrl(...))
  so the version request targets the backend port in desktop mode. The bare
  /api/health predates the Tauri webview (blame: the earlier onboarding commit,
  which ran with same-origin frontend/backend); in desktop mode the webview
  origin is tauri://localhost so the bare path fails.
- install.ps1: gate the install_python_stack.py hotfix on a sentinel comment
  instead of a content regex; append the sentinel after applying so reruns
  are unambiguous.
- unsloth_cli/commands/studio.py _write_auth_secret: use the atomic mkstemp +
  os.replace path on Windows too; chmod calls are wrapped in try/except OSError.
- studio/src-tauri/src/preflight.rs probe_existing_backends: fan out the health
  probes concurrently; desktop-auth status still runs sequentially per candidate.
  reqwest::Client is internally Arc-wrapped so the in-loop .clone() is a
  refcount bump, not a deep clone; annotated inline.
- studio/src-tauri/src/preflight.rs run_cli_probe: wait() after kill() to reap
  the child, matching probe_cli_capability.
- studio/src-tauri/src/process.rs + main.rs: add stop_backend_detached and use
  it from the tray quit handler so the 5s graceful-wait does not block the
  Tauri main loop. RunEvent::Exit keeps the synchronous safety-net call.
- studio/backend/main.py: drop the permissive localhost CORS regex in
  api-only mode; the explicit allow_origins list is sufficient.
- .github/workflows/release-desktop.yml: drop max-parallel: 1 so platform
  builds run in parallel, and lift releaseBody to an env var so the three
  tauri-action invocations share one source of truth.

* Fix review findings (loop 2)

- studio/backend/auth/storage.py update_password: clear_desktop_secret()
  alongside clear_bootstrap_password() so rotating the admin password
  also revokes any previously provisioned .desktop_secret. Without this,
  an old local desktop credential keeps minting fresh admin tokens via
  /api/auth/desktop-login after a password rotation.
- studio/src-tauri/src/desktop_auth.rs provision_desktop_auth: wrap
  cmd.output().await in tokio::time::timeout(30s). DESKTOP_AUTH_LOCK is
  held across the whole desktop_auth flow, and previously a hanging
  `unsloth studio provision-desktop-auth` subprocess would pin the lock
  indefinitely and freeze every subsequent desktop_auth call.

* Add review tests

* Consolidate review tests

Merge review-added tests into the existing studio/backend/tests/test_desktop_auth.py
(the PR's authoritative desktop-auth test file). Drops three scaffolding files under
tests/python/ in favor of five focused tests next to the tests they extend:
- test_update_password_clears_desktop_secret (runtime)
- test_update_password_on_unknown_user_leaves_desktop_secret_intact (runtime)
- test_cli_provisioning_delegates_to_storage_create_desktop_secret (source-level)
- test_cli_connect_auth_db_reads_storage_db_path (source-level)
- test_desktop_auth_provision_has_bounded_timeout (Rust source-level)

* Revert auth-guards.ts Tauri branches to unconditional form

The review loop on PR 5144 introduced a regression: the isTauri branch of
requireAuth redirected to /login when tauriAutoAuth() returned false, and
requireGuest / requirePasswordChangeFlow silently fell through on the same
condition. The Tauri desktop app authenticates via a local auto-generated
secret; it must never surface /login or /change-password to the user. A
failed auto-auth should let the startup layer retry, not expose a password
form.

Restore the three Tauri branches to the author's original unconditional
form (requireAuth: return; requireGuest / requirePasswordChangeFlow: throw
redirect({to: '/chat'})). Keep the rest of the review fixes -- the
apiUrl() fetch wrapping, authRedirect helper, and fetchAuthStatus refactor
are all legitimate improvements and are preserved.

* Revert release-desktop.yml to author's version

The review loop's workflow-file tweaks (drop max-parallel: 1, lift releaseBody
to an env var) are cosmetic. OAuth tokens cannot push workflow-file changes,
and fine-grained PATs cannot honor maintainerCanModify on a third-party fork.
Reverting the workflow file to wasimysaid's version lets the push go through
without needing a classic PAT with both repo and workflow scopes.

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

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

---------

Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Daniel Han <unslothai@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-04-23 04:50:10 -07:00
Datta Nimmaturi
9311df2b29
[Studio] multi gpu finetuning/inference via "balanced_low0/sequential" device_map (#4602)
* [WIP] balanced device map for studio

* gpus as a request parameter

* API for multi GPU stuff

* return multi gpu util in new API

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

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

* Use balanced_low0 instead of balanced

* Use balanced_low0 instead of balanced

* Fix device_map typo, UUID parsing crash, set() filter bug, and broken tests

- balanced_low0 -> balanced_low_0 (transformers/accelerate rejects the old string)
- get_parent_visible_gpu_ids() now handles UUID/MIG CUDA_VISIBLE_DEVICES
  gracefully instead of crashing on int() parse
- _get_backend_visible_gpu_info() set() or None bug: empty set is falsy so
  CUDA_VISIBLE_DEVICES=-1 would disable filtering and report all GPUs
- test_gpu_selection.py: add missing get_visible_gpu_utilization import and
  add required job_id arg to start_training() calls

* Smart GPU determinism using estimates

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

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

* disallow gpu selection for gguf for now

* cleanup

* Slightly larger baseline

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

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

* Treat empty list as auto

* Verbose logging/debug

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

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

* Cleanup and revert unnecessary deletions

* Cleanup excessive logs and guard against disk/cpu offload

* auth for visibility API. cleanup redundant imports. Adjust QLoRA estimate

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

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

* support for non cuda gpus

* Fix multi-GPU auto-selection memory accounting

The multi_gpu_factor was applied uniformly to all GPUs including the
first one, which unfairly penalizes single-GPU capacity when
transitioning to multi-GPU. This created a discontinuity where a model
that barely fits 1 GPU would suddenly require 2 GPUs because the first
GPU's free memory was discounted by 20%.

Now the first GPU keeps its full free memory, and only additional GPUs
have an overhead factor (0.85) applied to account for inter-GPU
communication and sharding overhead. This gives more accurate
auto-selection and avoids unnecessary multi-GPU for models that
comfortably fit on one device.

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

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

* Add sandbox tests for multi-GPU selection logic

24 tests covering model size estimation, memory requirements, automatic
GPU selection, device map generation, GPU ID validation, and multi-GPU
overhead accounting. All tests use mocks so they run without GPUs on
Linux, macOS, and Windows.

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

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

* Fix reviewer findings: 4bit inference estimate, fallback, GGUF gpu_ids, retry

1. 4-bit inference now uses reduced memory estimate (model_size/3 + buffer)
   instead of the FP16 1.3x multiplier. This prevents over-sharding
   quantized models across unnecessary GPUs.

2. When model size estimation fails, auto_select_gpu_ids now falls back to
   all visible GPUs instead of returning None (which could default to
   single-GPU loading for an unknown-size model).

3. GGUF inference route now treats gpu_ids=[] as auto-selection (same as
   None) instead of rejecting it as an unsupported explicit request.

4. Training retry path for "could not get source code" now preserves the
   gpu_ids parameter so the retry lands on the same GPUs.

5. Updated sandbox tests to cover the new 4-bit inference estimate branch.

* Remove accidentally added unsloth-zoo submodule

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

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

* Fix UUID/MIG visibility and update test expectations

1. nvidia.py: When CUDA_VISIBLE_DEVICES uses UUID/MIG tokens, the
   visibility APIs now return "unresolved" with empty device lists instead
   of exposing all physical GPUs. This prevents the UI from showing GPUs
   that the backend process cannot actually use.

2. test_gpu_selection.py: Updated test expectations to match the new
   multi-GPU overhead accounting (first GPU at full capacity, 0.85x for
   additional GPUs) and 4-bit inference memory estimation formula.
   All 60 tests now pass.

* Add CPU/disk offload guard to audio inference path

The audio model loading branch returned before the common
get_offloaded_device_map_entries() check, so audio models loaded with a
multi-GPU device_map that spilled layers to CPU/disk would be accepted
instead of rejected. Now audio loads also verify no modules are offloaded.

* Improve VRAM requirement estimates

* Replace balanced_low_0 with balanced

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

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

* refine calculations for slightly easier nums

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

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

* adjust estimates

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

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

* Use nums instead of obj to avoid seralisation error

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

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

* Harden nvidia-smi parsing and fix fallback GPU list

1. nvidia.py: Wrap int() casts for GPU index and memory in try/except
   so MIG slices, N/A values, or unexpected nvidia-smi output skip the
   unparseable row instead of aborting the entire GPU list.

2. nvidia.py: Handle GPU names containing commas by using the last
   field as memory instead of a fixed positional index.

3. hardware.py: fallback_all now uses gpu_candidates (GPUs with verified
   VRAM data) instead of raw devices list, which could include GPUs
   with null VRAM that were excluded from the ranking.

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

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

* cleanup

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

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

* consolidate raise_if_offload

* Improve MoE support. Guard against nvidia-smi failures

* Improve MoE support. Guard against nvidia-smi failures

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

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

* Fix shared-expert LoRA undercount, torch VRAM fallback, and apply_gpu_ids edge case

1. vram_estimation.py: compute_lora_params now includes shared experts
   (n_shared_experts) alongside routed experts when computing MoE LoRA
   adapter parameters. Previously only n_experts were counted, causing
   the estimator to undercount adapter, optimizer, and gradient memory
   for DeepSeek/GLM-style models with shared experts.

2. hardware.py: _torch_get_per_device_info now uses mem_get_info (which
   reports system-wide VRAM usage) instead of memory_allocated (which
   only reports this process's PyTorch allocations). This prevents
   auto-selection from treating a GPU as mostly free when another
   process is consuming VRAM. Falls back to memory_allocated when
   mem_get_info is unavailable.

3. hardware.py: apply_gpu_ids([]) now returns early instead of setting
   CUDA_VISIBLE_DEVICES="" which would disable CUDA entirely. Empty
   list inherits the parent visibility, same as None.

4. hardware.py: Upgraded fallback_all GPU selection log from debug to
   warning so operators are notified when the model likely will not fit
   in available VRAM.

* Guard nvidia-smi subprocess calls against OSError and TimeoutExpired

get_visible_gpu_utilization and get_backend_visible_gpu_info now catch
OSError (nvidia-smi not found) and TimeoutExpired internally instead
of relying on callers to wrap every invocation. Returns the standard
available=False sentinel on failure so the torch-based fallback in
hardware.py can take over.

* Guard get_primary_gpu_utilization and reset GPU caches between tests

1. nvidia.py: get_primary_gpu_utilization now catches OSError and
   TimeoutExpired internally, matching the pattern already used in
   get_visible_gpu_utilization and get_backend_visible_gpu_info. All
   three nvidia-smi callers are now self-contained.

2. test_gpu_selection.py: Added _GpuCacheResetMixin that resets the
   module-level _physical_gpu_count and _visible_gpu_count caches in
   tearDown. Applied to all test classes that exercise GPU selection,
   device map, or visibility functions. This prevents stale cache
   values from leaking between tests and causing flaky results on
   machines with real GPUs.

* Fix nvidia-smi fallback regression and physical GPU count validation

1. hardware.py: get_gpu_utilization, get_visible_gpu_utilization, and
   get_backend_visible_gpu_info now check result.get("available") before
   returning the nvidia-smi result. When nvidia-smi is unavailable or
   returns no data (e.g., containers without nvidia-smi, UUID/MIG masks),
   the functions fall through to the torch-based fallback instead of
   returning an empty result. This fixes a regression where the internal
   exception handling in nvidia.py prevented the caller's except block
   from triggering the fallback.

2. hardware.py: resolve_requested_gpu_ids now separates negative-ID
   validation from physical upper-bound validation. The physical count
   check is only enforced when it is plausibly a true physical count
   (i.e., higher than the largest parent-visible ID), since
   torch.cuda.device_count() under CUDA_VISIBLE_DEVICES returns the
   visible count, not the physical total. The parent-visible-set check
   remains authoritative in all cases. This prevents valid physical IDs
   like [2, 3] from being rejected as "out of range" when nvidia-smi is
   unavailable and CUDA_VISIBLE_DEVICES="2,3" makes torch report only
   2 devices.

* Fix UUID/MIG torch fallback to enumerate devices by ordinal

When CUDA_VISIBLE_DEVICES uses UUID or MIG identifiers,
get_parent_visible_gpu_ids() returns [] because the tokens are
non-numeric. The torch fallback in get_visible_gpu_utilization() and
get_backend_visible_gpu_info() previously passed that empty list to
_torch_get_per_device_info(), getting nothing back.

Now both functions detect the empty-list case and fall back to
enumerating torch-visible ordinals (0..device_count-1) with
index_kind="relative". This means the UI and auto-selection still
see real device data in Kubernetes, MIG, and Slurm-style UUID
environments where nvidia-smi output cannot be mapped to physical
indices.

Updated test_uuid_parent_visibility to verify the new torch fallback
path returns available=True with relative ordinals.

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

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

* Add type hint for gpu_ids parameter in InferenceOrchestrator.load_model

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
2026-03-30 02:33:15 -07:00