Offload the gated-base HEAD preflight to a worker thread
start_diffusion_training is async but called _preflight_gated_base inline; it does a blocking urllib urlopen HEAD to Hugging Face (up to a 5s timeout) to detect a gated/unauthorized base repo. On a slow or unreachable network that stalled the FastAPI event loop, freezing every concurrent status/progress/cancel request until it returned or timed out. Wrap it in asyncio.to_thread, matching the dataset preflight and GPU cleanup just below it. Regression test asserts it runs off the coroutine thread.
This commit is contained in:
parent
37185e651d
commit
eee5a53658
2 changed files with 36 additions and 1 deletions
|
|
@ -1301,7 +1301,12 @@ async def start_diffusion_training(
|
|||
# Preflight access to a gated base repo with the user's token BEFORE freeing GPU
|
||||
# residents, so a missing/insufficient token fails fast (400) without tearing down the
|
||||
# user's loaded chat/Images model, and never surfaces as a confusing mid-load 401.
|
||||
_preflight_gated_base(config.get("base_model", ""), config.get("hf_token"))
|
||||
# Offloaded to a worker thread: it does a blocking urlopen HEAD (up to a 5s timeout) to
|
||||
# Hugging Face, which would otherwise stall the event loop and every concurrent
|
||||
# status/progress/cancel request, as the filesystem preflight just below already does.
|
||||
await asyncio.to_thread(
|
||||
_preflight_gated_base, config.get("base_model", ""), config.get("hf_token")
|
||||
)
|
||||
|
||||
# Preflight the dataset too: a missing/empty/uncaptionable data_dir otherwise
|
||||
# fails inside the spawned trainer AFTER the user's chat/Images model was
|
||||
|
|
|
|||
|
|
@ -391,6 +391,36 @@ def test_route_start_frees_gpu_off_the_coroutine_thread(client, monkeypatch):
|
|||
assert threads["cleanup"] is not threads["inline"] # offloaded to a worker, not run inline
|
||||
|
||||
|
||||
def test_route_start_preflights_gated_base_off_the_coroutine_thread(client, monkeypatch):
|
||||
# _preflight_gated_base does a blocking urlopen HEAD (up to a 5s timeout) to Hugging Face, so
|
||||
# the async start route must offload it via asyncio.to_thread rather than run it inline and
|
||||
# freeze the event loop for concurrent status/progress/cancel requests. Assert it runs on a
|
||||
# DIFFERENT thread than the inline coroutine body (service.start), which an inline call could
|
||||
# not.
|
||||
import threading
|
||||
|
||||
import routes.training as tr
|
||||
|
||||
threads: dict = {}
|
||||
|
||||
def _record_preflight(base_model, hf_token):
|
||||
threads["preflight"] = threading.current_thread()
|
||||
|
||||
monkeypatch.setattr(tr, "_preflight_gated_base", _record_preflight)
|
||||
|
||||
orig_start = client._fake.start
|
||||
|
||||
def _record_start(config):
|
||||
threads["inline"] = threading.current_thread()
|
||||
return orig_start(config)
|
||||
|
||||
monkeypatch.setattr(client._fake, "start", _record_start)
|
||||
|
||||
r = client.post("/api/train/diffusion/start", json = _BODY)
|
||||
assert r.status_code == 200, r.text
|
||||
assert threads["preflight"] is not threads["inline"] # offloaded to a worker, not run inline
|
||||
|
||||
|
||||
def test_route_start_forwards_extra_training_knobs(client):
|
||||
# max_grad_norm and lora_target_modules must reach the service, not be silently dropped.
|
||||
body = {**_BODY, "max_grad_norm": 0.5, "lora_target_modules": ["to_q", "to_v"]}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue