9fc21b3977* Studio: make lifespan shutdown resilient to a dead default executor On an abrupt shutdown (closing the Windows console window, or interpreter teardown racing uvicorn's graceful stop) the event loop's default thread-pool executor can already be shut down by the time the FastAPI lifespan shutdown runs. The first post-yield statement was an unguarded `await asyncio.to_thread(terminate_hub_downloads)`, so executor.submit raised `RuntimeError: cannot schedule new futures after shutdown`. That raise propagated up through every nested merged_lifespan __aexit__, aborted the rest of the cleanup (DEVICE reset, compiled-cache clear), and surfaced as "Application shutdown failed. Exiting." Extract the post-yield cleanup into utils/lifespan_shutdown.run_lifespan_shutdown and guard each step independently. On the to_thread RuntimeError, fall back to running the (already best-effort, quick) terminate inline on the loop thread so shutdown still completes cleanly. The helper is dependency-injected and free of the heavy backend import graph, so it is unit-tested in isolation. Add tests/test_lifespan_shutdown.py (4 cases: dead-executor survival, normal path, terminate error, clear error). Validated on windows-latest and ubuntu-latest runners on Python 3.13. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Studio: only retry terminate inline when scheduling fails, not when the body raises Address review feedback on the shutdown helper: the previous `except RuntimeError` after `asyncio.to_thread(terminate_downloads)` could not tell a dead-executor scheduling failure from a RuntimeError raised by terminate_downloads itself, so a body-side RuntimeError on a healthy executor ran the cleanup a second time inline. Schedule via loop.run_in_executor and await separately: a dead default executor raises synchronously at submit time (inline fallback), while a body exception only surfaces when the future is awaited (logged, never retried). Add a regression test that a body RuntimeError runs terminate exactly once. * Studio: run terminate cleanup with a copied context (parity with asyncio.to_thread) Simulation across the executor-state x exception x DEVICE matrix surfaced the one behavioural difference from the original implementation: asyncio.to_thread copies the caller's contextvars into the worker thread, while a bare run_in_executor(None, fn) does not. Restore exact parity by scheduling ctx.run(terminate_downloads) from a contextvars.copy_context(), so the refactor is a behavioural no-op apart from the intended dead-executor recovery. Add a regression test asserting the copied context is visible to terminate_downloads. * Studio: tighten comments in lifespan_shutdown helper and tests Condense the verbose docstrings/comments to the non-obvious rationale and drop the self-evident ones. Verified comment-only with comment_tools.py check --strip-docstrings (code signature unchanged); tests and sims still green. * Studio: address review nits on lifespan shutdown helper Annotate hw_module as types.ModuleType, reword the schedule/await comment (inline fallback runs, it does not retry), and add a test for the public loop.shutdown_default_executor() path (the 'Executor shutdown has been called' RuntimeError that real uvicorn shutdown takes, distinct from the submit-time 'cannot schedule new futures after shutdown'). --------- Co-authored-by: Daniel Han <michaelhan2050@gmail.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
134 lines
4.2 KiB
Python
134 lines
4.2 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
"""Regression tests for run_lifespan_shutdown: a dead default executor (the
|
|
abrupt-shutdown teardown race) must not abort the remaining cleanup. The helper
|
|
is dependency-injected, so these need only structlog."""
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import types
|
|
|
|
from utils.lifespan_shutdown import run_lifespan_shutdown
|
|
|
|
|
|
def _counter():
|
|
box = {"n": 0}
|
|
|
|
def _fn():
|
|
box["n"] += 1
|
|
|
|
return box, _fn
|
|
|
|
|
|
def test_run_lifespan_shutdown_survives_dead_default_executor():
|
|
term_box, terminate = _counter()
|
|
clear_box, clear = _counter()
|
|
hw = types.SimpleNamespace(DEVICE = "cuda:0")
|
|
|
|
async def _drive():
|
|
loop = asyncio.get_running_loop()
|
|
# Kill the default executor to mimic the teardown race.
|
|
await asyncio.to_thread(lambda: None)
|
|
loop._default_executor.shutdown(wait = True)
|
|
await run_lifespan_shutdown(terminate, clear, hw)
|
|
|
|
asyncio.run(_drive())
|
|
|
|
assert term_box["n"] == 1, "terminate must run via inline fallback"
|
|
assert clear_box["n"] == 1, "clear must still run after the to_thread failure"
|
|
assert hw.DEVICE is None
|
|
|
|
|
|
def test_run_lifespan_shutdown_survives_shutdown_default_executor():
|
|
"""Production path: loop.shutdown_default_executor() makes run_in_executor raise
|
|
'Executor shutdown has been called'; the helper must still recover inline."""
|
|
term_box, terminate = _counter()
|
|
clear_box, clear = _counter()
|
|
hw = types.SimpleNamespace(DEVICE = "cuda:0")
|
|
|
|
async def _drive():
|
|
await asyncio.get_running_loop().shutdown_default_executor()
|
|
await run_lifespan_shutdown(terminate, clear, hw)
|
|
|
|
asyncio.run(_drive())
|
|
|
|
assert term_box["n"] == 1, "terminate must run via inline fallback"
|
|
assert clear_box["n"] == 1
|
|
assert hw.DEVICE is None
|
|
|
|
|
|
def test_run_lifespan_shutdown_normal_path():
|
|
"""Healthy executor: each step runs exactly once."""
|
|
term_box, terminate = _counter()
|
|
clear_box, clear = _counter()
|
|
hw = types.SimpleNamespace(DEVICE = "cuda:0")
|
|
|
|
asyncio.run(run_lifespan_shutdown(terminate, clear, hw))
|
|
|
|
assert term_box["n"] == 1
|
|
assert clear_box["n"] == 1
|
|
assert hw.DEVICE is None
|
|
|
|
|
|
def test_run_lifespan_shutdown_swallows_terminate_errors():
|
|
"""A terminate failure must not block later cleanup."""
|
|
clear_box, clear = _counter()
|
|
hw = types.SimpleNamespace(DEVICE = "cuda:0")
|
|
|
|
def _boom():
|
|
raise ValueError("boom")
|
|
|
|
asyncio.run(run_lifespan_shutdown(_boom, clear, hw))
|
|
|
|
assert clear_box["n"] == 1, "later cleanup must run even when terminate raises"
|
|
assert hw.DEVICE is None
|
|
|
|
|
|
def test_run_lifespan_shutdown_swallows_clear_errors():
|
|
"""A clear failure must not raise out of shutdown."""
|
|
term_box, terminate = _counter()
|
|
hw = types.SimpleNamespace(DEVICE = "cuda:0")
|
|
|
|
def _boom():
|
|
raise ValueError("boom")
|
|
|
|
asyncio.run(run_lifespan_shutdown(terminate, _boom, hw))
|
|
|
|
assert term_box["n"] == 1
|
|
assert hw.DEVICE is None
|
|
|
|
|
|
def test_run_lifespan_shutdown_does_not_retry_body_runtime_error():
|
|
"""A body-side RuntimeError (healthy executor) must run terminate once, not retry inline."""
|
|
term_box, _ = _counter()
|
|
clear_box, clear = _counter()
|
|
hw = types.SimpleNamespace(DEVICE = "cuda:0")
|
|
|
|
def _boom():
|
|
term_box["n"] += 1
|
|
raise RuntimeError("body failed")
|
|
|
|
asyncio.run(run_lifespan_shutdown(_boom, clear, hw))
|
|
|
|
assert term_box["n"] == 1, "body RuntimeError must not be retried inline"
|
|
assert clear_box["n"] == 1, "later cleanup must still run"
|
|
assert hw.DEVICE is None
|
|
|
|
|
|
def test_run_lifespan_shutdown_preserves_contextvars():
|
|
"""terminate runs in a copy of the caller's context (parity with asyncio.to_thread)."""
|
|
cv = contextvars.ContextVar("unsloth_test_cv")
|
|
cv.set("bound-value")
|
|
seen = []
|
|
clear_box, clear = _counter()
|
|
hw = types.SimpleNamespace(DEVICE = "cuda:0")
|
|
|
|
def terminate():
|
|
seen.append(cv.get("UNSET"))
|
|
|
|
asyncio.run(run_lifespan_shutdown(terminate, clear, hw))
|
|
|
|
assert seen == ["bound-value"], "terminate must run with the caller's contextvars"
|
|
assert clear_box["n"] == 1
|
|
assert hw.DEVICE is None
|