Studio: Auto download transport, Xet memory caps, and stall detection on the hub path - #7742
Studio: Auto download transport, Xet memory caps, and stall detection on the hub path#7742danielhanchen wants to merge 26 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Verified against a live Studio, not mocksA Studio was installed from this branch (
The demotion round trip is the interesting one, because it exercises the whole chain rather than any single piece: the zoo ladder records a failed Xet attempt, the verdict persists next to the HF cache, the backend's capabilities endpoint reports
Anonymous CI, no HF_TOKENReplicated onto a staging repo and run on
56% less peak RSS for 12-30% more wall time. Anonymous Xet is fast and works, so no token is required for any of this. That CI job also caught a real bug that review did not: The multi-shard caseOn a 61GB / 13-shard safetensors repo ( |
UI and training, confirmedPlaywright screenshot of the Model hub header on this branch shows the picker rendering Auto | HTTP | Xet with Auto selected, and the tooltip explaining itself rather than being an opaque third option:
When the machine has been demoted the same tooltip reads Training was exercised as a real run rather than a status ping, because the point is the training worker's own model download going through the shared ladder: So all three surfaces named in the issue are covered end to end on a live server: hub (download under Auto, completed), inference (GGUF load + chat returning One caveat worth statingThe hub stall path is covered by unit tests (kill -> error -> existing HTTP retry, plus the already-exited race and the degraded-unsloth_zoo case), but I did not manufacture a live hang through the hub worker specifically -- |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1e3945fd53
ℹ️ 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".
| if stalled: | ||
| # A machine whose Xet transfers hang is one that should stop starting on Xet. | ||
| _record_xet_failure(stalled[0], logger) |
There was a problem hiding this comment.
When a Xet worker exits with a retryable CAS/network error before the stall watchdog fires, state == "error" and the HTTP retry below runs, but stalled remains empty so the health tracker never sees that failed Xet attempt. Since Auto's backend verdict is based on recorded outcomes, subsequent Auto downloads on the same machine keep choosing Xet and paying the same failed attempt before falling back; record a Xet failure for every Xet-to-HTTP retry, not only watchdog stalls.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
The recording is deliberately scoped to stalls. A fast CAS or 5xx error costs one quick HTTP retry, whereas a hang costs the full stall timeout, and counting transient blips would demote machines on failures HTTP would have hit too. That noise is what the two-strike threshold exists to absorb.
| logger.exception("%s failed to kill stalled worker for %s", log_prefix, label) | ||
|
|
||
| try: | ||
| return start_watchdog( |
There was a problem hiding this comment.
This new hub watchdog call relies on the shared start_watchdog defaults, so the hub path ignores the Studio-side DEFAULT_STALL_TIMEOUT = 30.0 added in this commit. In installs where the shared watchdog default is still 180s, a frozen Xet hub download sits with no progress for three minutes before the HTTP retry instead of using the intended 30s Xet fallback; pass the explicit stall timeout here so hub downloads get the new deadline.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Passing an explicit value here would break the intended behaviour: the shared layer resolves None through UNSLOTH_XET_STALL_TIMEOUT before falling back to its per-transport default, so hardcoding 30s would silently disable that operator override. The 180s you describe comes from the older shared release, which the dependency floor bump addresses rather than this call site. Even at 180s the hub path now recovers, where before this PR it never did.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
Measured results126 anonymous cells over six staging-CI runs ( The memory problem is real and unbounded
15.38 GB on a 16 GB runner, measured across the process tree. The variance matters more than the
Stock swings by a factor of 250 on one file; capped stays inside a factor of 2 everywhere. Cost of Anonymous users should keep XetCapped Xet is 4.3x faster than HTTP with no token (1728 vs 402 Mbps median). A separate local On a slow link Xet's advantage disappearsThrottled through a userspace CONNECT proxy (only the download process is slowed; the host link is
All three saturate the link. Xet only wins where the link is not the bottleneck, and peak RSS stays Detection, live
Two of the seven repos ship a single 17 GB / 22 GB file, larger than a runner's disk. Those CI cells |
…stalled hub download The model hub is how most people download, and it was the one download path with no stall detection at all. finalize_worker_exit says so explicitly: it relies on the worker's exit code, and a Xet transfer that hangs with no progress and no error never produces one. That is the frozen progress bar users see. The inference and training paths already ran the unsloth_zoo watchdog; the hub now runs it too, and a stalled worker is SIGKILLed so it lands as "error" -- the exact state the existing XET-to-HTTP retry already keys on. Nothing about that retry path changes. Downloads spawned by the hub now get RAM-derived HF_XET_* caps. Note this cannot be done with setdefault alone: HF_XET_HIGH_PERFORMANCE is inherited from the parent environment, and xet-core applies that preset AFTER reading the environment, so leaving it on discards every cap rather than competing with it. It is cleared explicitly, with UNSLOTH_XET_ALLOW_HIGH_PERFORMANCE to opt back in. The transport picker gains a third option, Auto, which becomes the default. Auto resolves server-side because only the backend can see this machine's RAM, its hf_xet build, and whether Xet has been failing here. It resolves through the capabilities endpoint rather than independently on both sides: the resolved transport is compared against the .transport marker of any existing partial, so client and server must agree or a resume would be misjudged. "auto" is deliberately not a member of VALID_TRANSPORTS -- the on-disk marker must keep naming the writer that produced a partial. An explicit Xet or HTTP choice is still honoured, and a previously stored preference is preserved. The tooltip reports what Auto currently resolves to and why. poll-loop carried the stored preference as a boolean (`getTransportMode() === XET`), which would have read "auto" as "not xet" and sent every download over HTTP; it now stays unresolved until effectiveTransportMode().
for more information, see https://pre-commit.ci
unsloth_zoo.__init__ runs torch accelerator detection and raises NotImplementedError on a CPU-only host. _load_optional gave up there, which silently disabled the RAM caps on exactly the small machines they exist to protect. _load_shared already had this retry; _load_optional now does too, with the flag scoped so it cannot leak into unrelated later imports. Caught by the staging CI download job on an ubuntu-latest runner.
for more information, see https://pre-commit.ci
…g module CI on a runner with the released unsloth_zoo installed caught two things review did not. The high-performance flag was cleared through xet_env_overrides(), which returns nothing when the installed zoo predates hf_xet_tuning -- and that same older zoo is the one setting HF_XET_HIGH_PERFORMANCE=1 at import. So on exactly the installs that cannot be fixed by upgrading Studio alone, the worker inherited a 64GB buffer ceiling. Clearing it is now unconditional. The two cap assertions now skip when the tuning module is genuinely absent, since there are no caps to assert, and a new test pins the clear-without-tuning case. Also names TRANSPORT_AUTO in the cache layer's error path: 'auto' is a request preference that only the server can resolve, so reaching prepare_cache_for_transport with it is a specific bug worth a specific message rather than a generic 'invalid transport'.
for more information, see https://pre-commit.ci
The verdict is deliberately sticky across sessions: two consecutive Xet failures pin a machine to HTTP for 24h. That also means any machine which has genuinely had a Xet stall starts the ladder on HTTP inside tests that expect it to start on Xet, and test_shim_injects_studio_prepare_on_http_retry then sees the fallback happen without the Xet attempt it asserts. Reproduced locally against a real demoted verdict; a clean CI runner has no state file so it hides there. Mirrors the fixture already in the unsloth_zoo suite. Verified no collateral damage from the HF_HOME redirect: 1412 passed across the hub / cache / download / transport tests.
1100754 to
9fa856e
Compare
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
The red
|
| main commit | result |
|---|---|
aa4ea75 (this PR's original base) |
success |
3044401 (current main, the rebase target) |
failure, identical assertions |
Between those sit 63307b1 (#7736, default Show all quantizations off, collapse single-quant rows),
6c60158 (#7745, row actions on collapsed quant rows), 13fdd06 (#7747, remember the Run Settings
advanced toggle) and 3044401 (#7755). Those change exactly the surface the check drives -- opening
run-settings from a quant row, and the Reset button inside it -- and the Playwright script was not
updated with them.
Matching behaviour on this branch: based on aa4ea75 it failed once and passed on re-run (a real
flake); rebased onto 3044401 it fails consistently, 2 for 2. Mac Studio UI CI and Windows Unsloth
UI CI both pass on the same commit, so it is Linux-only, and the frontend diff here is confined to
features/hub/**, which nothing in Chat run-settings imports.
Leaving it alone rather than patching a model-picker regression inside a Xet transport PR.
Everything else is green: 39 checks including the full Python matrix, all three Core combos,
Repo tests (CPU), Source lint, CodeQL and every Windows/macOS/Linux job.
…watchdog, blocking probe Carry the Auto verdict through capability normalization. The normalizer rebuilt the object from http and xet alone, so auto_resolves_to and auto_reason were dropped on every response: Auto resolved to Xet on every machine including ones the backend had just demoted to HTTP, and the toggle always read "Xet". The whole server-side verdict was computed, serialized and thrown away in the browser. The normalizer now lives in transport-capabilities.ts, because api.ts imports the auth barrel and cannot be loaded by the test runner, which is exactly how this shipped green; three tests now cover it. Stop the lifecycle tests from running the watchdog inline. _ImmediateThread is installed on the stdlib threading module, which the shared Zoo watchdog imports too, so the watchdog ran synchronously and blocked in Event.wait() before finalize_worker_exit could set its stop flag. hub/tests/test_download_lifecycle.py hung forever; it now finishes in 0.13s. CI never caught this because studio-backend-ci runs tests/, not hub/tests/. Scope the hub watchdog to its own worker. child_pid is only read under watch_new_partials_only, so the measurement stayed repo-wide and two same-transport GGUF variants of one repo (which the registry deliberately allows to run concurrently) reset each other's stall timer, leaving a hung variant with no fallback. Record a completed Xet download, not just a stall, so "two failures in a row" means in a row. A stall today and another next week were counted as consecutive despite every download in between succeeding, pinning Auto to HTTP for 24h. Omit unset timeout kwargs in the shim rather than forwarding None. No production caller passes them, and an older unsloth_zoo takes None literally: its watchdog hands interval to Event.wait(), which then blocks forever, so a hung download would never fall back. Resolve the transport off the event-loop thread. Both async handlers called the resolver synchronously one line after correctly awaiting to_thread for the repo-id lookup, and resolving "auto" can run the Xet reachability probe, whose 3s budget does not cover DNS. Give the capabilities cache a 30s TTL. It was held for the page session, so after a mid-session demotion Auto kept sending an explicit "xet" (which deliberately bypasses health) until reload.
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…t override Let Auto actually use the reachability probe. The client resolves Auto from the capabilities endpoint, which deliberately answers without probing because the UI polls it on render, and then sends a concrete transport. The probe this PR added therefore never ran for a real user: a host whose CAS is unreachable but which has not recorded a failure yet discovered that by stalling for 30s rather than at preflight. The endpoint now takes an opt-in probe flag, and only effectiveTransportMode sets it, since that runs at download start rather than on render. Resolution stays client-side so the chosen transport still matches the .transport marker of any partial, and the cache records whether its entry was probed so a cheap answer cannot satisfy a probing caller. Serialize the GPU-init override. The retry that imports an optional helper with UNSLOTH_ZOO_DISABLE_GPU_INIT=1 saves, sets and restores process-wide state, and on the currently declared zoo floor those modules are absent, so it runs on every download. Two concurrent requests could interleave save/set/restore and leave the variable set for the life of the process, after which every worker inherits it and skips Zoo's GPU init. A lock around the sequence is enough; I deliberately did not memoize the result, since that would stop a newly installed zoo being picked up and would change what the existing loader tests observe.
for more information, see https://pre-commit.ci
|
@codex review |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: f465e47065
ℹ️ 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".
| # high-performance preset AFTER reading the environment). Overwrite, never setdefault. | ||
| for key in ("HF_XET_HIGH_PERFORMANCE", "HF_XET_HP"): | ||
| env[key] = "0" | ||
| for key, value in xet_env_overrides().items(): |
There was a problem hiding this comment.
On a cold backend when the background warm-up is disabled, unfinished, or skipped on a CPU-only host, an explicit or legacy Xet request reaches this call without the Auto health probe having loaded Zoo first. Both async download handlers invoke launch_worker synchronously, and xet_env_overrides() can import unsloth_zoo, Transformers, and Torch and then retry GPU initialization, blocking the event loop and freezing every Studio request for the duration. Resolve and cache these overrides in a worker thread before launching the process.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
spawn_worker can only be reached with use_xet=True after resolve_requested_use_xet, which both handlers already run under asyncio.to_thread and which resolves availability through xet_health() -- that is the call importing unsloth_zoo, transformers and torch and running the GPU-init retry. By the time xet_env_overrides() runs on the loop only unsloth_zoo.hf_xet_tuning is left, which is stdlib-only: measured 4.06s for the health import in the thread versus 0.000s for the tuning import afterwards.
…the real HF cache in tests Both of these are corrections to fixes from the previous round. The cached-job gate measured repo-wide bytes, but the registry deliberately lets two same-transport GGUF variants of one repo run concurrently and they share one blobs/ dir, so a cached no-op worker could be credited with its sibling's bytes. That does more than clear a streak: recording a Xet success also flips an already demoted verdict back to Xet. Variant jobs are now measured against their own blob hashes, and a variant whose hashes did not resolve is unmeasurable and never clears the streak. Job shapes that cannot have a concurrent same-repo sibling keep the repo-wide measure. Moving HF_HOME to session scope fixed the fixture ordering but widened the blast radius: HF_HOME also defaults HF_HUB_CACHE, HF_XET_CACHE and HF_TOKEN_PATH, and unlike the function-scoped version this one does reach the spawned E2E server. That sent it to an empty cache and an empty token store, i.e. a redownload of the default GGUF inside the 120s startup deadline and no credentials for a private --unsloth-model. Those three are now pinned to what the hub resolved from the real environment before HF_HOME moves, so only the Xet health file is isolated.
for more information, see https://pre-commit.ci
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…byte trips out of health Found by adversarial self-review, and the first item is a regression the sibling zoo change made live rather than a latent one. The hub arms the watchdog with no connect_timeout, so it inherited the shared 90s zero-byte default. That default is sized for a single-file download whose pre-byte phase is one HEAD. This worker calls snapshot_download(max_workers=1), so its pre-byte phase is a model_info lookup with retries plus one sequential HEAD per file, and for an already cached repo that is the entire job with no byte ever written. A few hundred files on a slow link exceeds 90s legitimately, so a healthy worker was killed before it started. Now 600s. Independently, do not let a pre-byte trip poison the machine's health record. Two recorded failures pin it to HTTP for 24h, and a connect-phase trip means no byte ever arrived, which is as likely to be slow metadata, a long queue of HEADs, or a cache lock as a broken Xet. The HTTP retry still happens either way, so this costs nothing and removes the expensive direction of error. Verified with a negative control that the new test fails when the branch is removed. The per-test Xet health isolation silently did nothing on CPU-only hosts: a bare `from unsloth_zoo import hf_xet_health` raises NotImplementedError there because unsloth_zoo's __init__ runs accelerator detection, and the fixture swallowed it as "degraded unsloth_zoo, nothing to isolate" when the module was in fact perfectly available. It now loads through the same _load_optional the shim uses, whose GPU-init retry exists for exactly those hosts. Two comments corrected to match their code: the variant watchdog scopes the data clock only (the pre-byte phase is still repo-wide by design, which is what spares a lock wait behind a live sibling), and the probe note contrasted itself with the wrong caller.
for more information, see https://pre-commit.ci
A small download can write and even finalize its blobs while the parent is still registering the process, so a baseline taken inside register_worker races it: bytes_after is then not greater than the baseline, a real Xet transfer is misread as a no-op, and the failure streak is never cleared. Two stalls either side of that success then read as consecutive and demote Auto to HTTP for 24h, which is the outcome the success-recording exists to prevent. launch_worker now samples before spawn() and passes the value in; register_worker keeps its own sampling for callers that do not.
for more information, see https://pre-commit.ci
This comment was marked as resolved.
This comment was marked as resolved.
Sorry, something went wrong.
…e on HTTP starts The connect_timeout I added last round is not accepted by the supported floor. unsloth_zoo 2026.8.1 declares start_watchdog keyword-only with no **kwargs and no connect_timeout, so the call raised TypeError, the surrounding except Exception swallowed it, and _start_stall_watchdog returned None. On every supported install that meant no watchdog thread at all: a hung Xet worker was never killed and the HTTP retry, which is triggered by that kill landing as an error, never fired. The feature was entirely off, not degraded. Verified by calling the real 2026.8.1 wheel: TypeError before, a live Event after. Fixed in the shim rather than at the call site, because this is the third version-skew bug in this PR and the first two were each fixed one symptom at a time. start_watchdog now drops any kwarg the installed callee does not accept, passing everything through when it takes **kwargs or cannot be introspected, so a newer zoo still gets the newer knobs and the next new kwarg is a no-op instead of a repeat of this. Dropping the pre-byte budget on 2026.8.1 is safe: that release resets its timer whenever the child owns no .incomplete, which is exactly the connect phase. An audit of the whole call site found connect_timeout to be the only unsupported kwarg today, and heartbeat_interval the only other one the unreleased zoo adds. Separately, only sample the Xet byte baseline for Xet starts. It is consumed solely by the Xet success-recording, and sampling it lazy-loads unsloth_zoo, and so torch and transformers, on the request path. Correction to the report on this one: the commit it cites moved that call rather than adding it, since register_worker already sampled unconditionally on the same synchronous path.
for more information, see https://pre-commit.ci
|
Codex Review: Didn't find any major issues. Chef's kiss. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
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 allow-list I added last round keyed on the "no progress" wording, which only covers the branch that fires while an .incomplete is growing. The hub worker runs snapshot_download(max_workers=1), so between two files it owns no partial at all and a hang there trips "did not resume" instead. That is the shape this worker hangs in most often, and the allow-list was silently dropping it, leaving Auto's health record blind to the most likely real failure. Bytes have already flowed by the time that verdict fires, so it is real evidence and belongs in the record. Only the pre-byte trip stays excluded: "did not start" means not one byte ever arrived, which a cold cache or a slow token handshake produces on a perfectly healthy machine. Also require state == "error". The watchdog appends its verdict from its own thread, so a worker that completed or was cancelled in the same instant would otherwise be charged a failure it did not earn.
…ildren The Xet shim sets UNSLOTH_ZOO_DISABLE_GPU_INIT=1 process-wide while it retries an optional unsloth_zoo import on a GPU-less host, then restores it. A worker spawned by another thread inside that window inherits it for its entire life, and unsloth_zoo answers that flag by installing stub triton and bitsandbytes modules, so the child would silently train against no-ops with nothing in the logs to say why. utf8_child_env now drops the variable while a loader is holding it. The loader only counts itself as the owner when it actually introduced the value, so an operator who exported the flag deliberately still has it reach their children.
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: dd043e8329
ℹ️ 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".
| """Copy *env* (or the current environment) with UTF-8 stdio forced.""" | ||
| child = dict(os.environ if env is None else env) | ||
| child["PYTHONIOENCODING"] = "utf-8" | ||
| if env is None and child.get("UNSLOTH_ZOO_DISABLE_GPU_INIT") == "1": |
There was a problem hiding this comment.
When the Xet loader's temporary UNSLOTH_ZOO_DISABLE_GPU_INIT=1 window overlaps a subprocess spawn that supplies an env mapping, this guard is skipped because env is None is false. Several training/probe spawns build an explicit environment by first copying os.environ (for example via child_env_without_native_path_secret() before utf8_child_env(...)), so they can still inherit the loader's transient flag for their whole lifetime and run with unsloth_zoo's triton/bitsandbytes stubs. Since gpu_init_override_active() already distinguishes the loader-owned value from an operator-set value, apply the scrub to explicit envs too.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
You are right that child_env_without_native_path_secret copies os.environ, so the flag does survive into those dicts. I walked every caller that passes an explicit env and none of the children is harmed by it: the two transformers sidecar probes, the vision check, the Vulkan probe, the torch version probe, the uv and pip installs, and the llama.cpp and whisper.cpp installers. The only one that imports unsloth_zoo at all is the DiffusionGemma shim, and I checked it directly, neither shim.py nor visual_engine.py references triton, bitsandbytes or any of the GPU init constants.
The children that would genuinely run against the stubs are the multiprocessing spawn workers, and those take no env argument at all, so filtering a dict could never have reached them. That hole is closed in 6ceeab2 by holding the loader lock across Process.start() in child_environment_for_spawn, with 5f3e783 making the lock reentrant since that context manager nests. Leaving the explicit env path as is.
| return (False, "hf_xet is not installed") | ||
| try: | ||
| from utils.hf_xet_fallback import xet_health | ||
| health = xet_health() |
There was a problem hiding this comment.
For clients that use the new backend transport_mode: "auto" field directly, this call reads the health verdict without opting into a probe, even though the download-start path is the moment intended to pay that cost. On a host with hf_xet installed, no cached health verdict, and an unreachable CAS endpoint, Auto will still choose the optimistic Xet path and only fall back after the stall watchdog fires; pass probe=True here so backend-resolved Auto avoids the failed Xet attempt up front.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
This is already doing what you are asking for. xet_health in the companion zoo is declared def xet_health(*, force: bool = False, probe: bool = True), so the bare call in resolve_auto_use_xet does run the reachability probe; it is the capabilities endpoint that explicitly passes probe=False, because the UI polls that one on render. The split is deliberate and the docstring right above the call says so.
So on a host with hf_xet installed, no cached verdict and an unreachable CAS, backend Auto resolves to HTTP up front rather than paying a stalled Xet attempt. Both download callers already run the resolution through asyncio.to_thread, with the 3s probe budget, so it never sits on the event loop even when DNS is blackholed.
My last commit guarded utf8_child_env, which is the wrong door. Training and inference workers are started with multiprocessing spawn (_CTX.Process in core/training/training.py and core/inference/orchestrator.py), and a spawn child copies the parent's live os.environ; there is no env dict to filter. The only four utf8_child_env callers that pass no env are pip invocations, which never import unsloth_zoo, so the hazard the commit described was left entirely open. child_environment_for_spawn already holds a lock across Process.start(), so it now also holds the shim's loader lock. A spawn cannot begin while a loader has UNSLOTH_ZOO_DISABLE_GPU_INIT set for its own import, which is the property that keeps stub triton and bitsandbytes out of a worker that would otherwise train against no ops with nothing in the log to say why. The loaders never spawn, so there is no lock order cycle, and the test asserts the shared lock structurally rather than trying to hit a microsecond window by timing. _load_optional now memoises its result, including the failure. On a zoo that predates these modules the import can never start succeeding, and without memoisation every xet_health, record_xet_outcome and xet_env_overrides call re-ran the whole GPU-init retry, re-opening that process-wide window on every single download. That also makes the new barrier uncontended: the window opens at most once per module per process. The ownership counter is also claimed before the env write and released after the restore, so the window in which the variable is set sits strictly inside the window in which it can be seen to be ours.
|
@codex review |
for more information, see https://pre-commit.ci
child_environment_for_spawn nests, which is why its own _spawn_env_lock is an RLock. Holding a plain Lock alongside it made the inner enter block forever, and the repo's existing test_nested_child_environment_for_spawn hung instead of failing. An RLock still excludes other threads, which is all the env save, set and restore sequence needs. The new test bounds its own wait, so the next person to get this wrong sees a failure rather than a suite that never finishes.
1 participant
Why
Three separate problems with Xet downloads in Studio.
The model hub had no stall detection at all.
finalize_worker_exitsays so in its own docstring: it relies on the worker's exit code, and a Xet transfer that hangs with no progress and no error never produces one. Inference (core/inference/worker.py) and training (core/training/worker.py) already ran the shared watchdog; the hub -- the path most users actually download through -- did not. That is the frozen progress bar.Nothing capped Xet's memory.
hf_xetsizes its reconstruction buffers from constants: 2GB floor + 512MB per concurrent file (8 of them), capped at 8GB, plus a 1GB prefetch floor. WithHF_XET_HIGH_PERFORMANCEinherited from the environment that cap becomes 64GB and the stream count 124.Nothing remembered.
resolve_effective_use_xetonly downgraded whenhf_xetwas missing, so a machine that had failed Xet ten times still started on Xet, and paid the stalled attempt, every single time.What changed
Stall detection on the hub path.
register_workernow starts the sharedstart_watchdogfor Xet jobs and SIGKILLs a worker that stops making byte-level progress. The kill is deliberate: the worker traps SIGTERM and exits 130, whichclassify_exitreads as a user cancel and would skip the retry; an untrapped kill lands aserror, which is the exact state the existing XET-to-HTTP retry already keys on. That retry path is unchanged. The worker's writer is sequential and resumable, so the partial is not lost work. The watchdog is stopped as soon as the worker is reaped, so post-download symlinking and verification are never read as a stall.Memory caps on spawned workers.
spawn_workermerges RAM-derivedHF_XET_*values into the worker env. This cannot be done withsetdefaultalone:envis seeded from the parent environment, so an inheritedHF_XET_HIGH_PERFORMANCE=1arrives here, and xet-core applies that preset after reading the environment -- it would discard every cap rather than compete with it. It is overwritten explicitly, withUNSLOTH_XET_ALLOW_HIGH_PERFORMANCE=1to opt back in. Operator-setHF_XET_*values are still preserved.An Auto transport, defaulted. The picker gains a third option. Auto resolves server-side, because only the backend can see this machine's RAM, its
hf_xetbuild, and whether Xet has been failing here.It resolves through the capabilities endpoint (
auto_resolves_to/auto_reason) rather than independently on both sides. That matters: the resolved transport is compared against the.transportmarker of an existing partial, so if the client picked Xet and the server picked HTTP a resume would be misjudged. One source of truth, resolved once.autois deliberately not a member ofVALID_TRANSPORTS. The on-disk marker must keep naming the writer that produced a partial;VALID_TRANSPORT_MODESis the separate set for request preferences.An explicit Xet or HTTP choice is still honoured (a user who asks for Xet gets Xet -- with the caps and the fallback, but the health verdict does not overrule them), and a preference already in localStorage is preserved. The tooltip says what Auto currently resolves to and why, e.g. "Currently: HTTP (Xet failed 2 times in a row on this machine)".
poll-loop.tscarried the stored preference as a boolean (getTransportMode() === TRANSPORT.XET), which would have read"auto"as "not xet" and sent every download over HTTP. It now stays unresolved untileffectiveTransportMode().Measured
HTTP is 15-30x slower than Xet on this link, which is why Auto keeps Xet first and treats HTTP strictly as recovery rather than as a safer default. Anonymous downloads (no
HF_TOKEN) reach Xet at full speed, so none of this depends on having a token.Tests
studio/backend/tests/test_hub_download_transport_auto.py(19 new): transport selection including "explicit Xet beats an unhealthy verdict" and the legacyuse_xetpath; the caps reaching the worker env including the high-performance override; stall -> kill -> retry including an already-exited worker and a degraded unsloth_zoo.studio/frontend/tests/transport-mode-auto.test.ts(4 new) pins the Auto default and that an existing explicit preference survives.Existing
test_hf_xet_fallback.py,test_gguf_xet_fallback_integration.pyandtest_hf_cache_settings.pypass unchanged; frontendnpm run typecheckis clean.Depends on unslothai/unsloth-zoo#972 for
hf_xet_tuning/hf_xet_health. Both are imported through the existing lazy, degrade-don't-crash shim inutils/hf_xet_fallback.py, so an older unsloth_zoo keeps working -- it just loses the caps and the verdict, not the ability to download.