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>
56 lines
2 KiB
Python
56 lines
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
|
|
|
|
"""Resilient FastAPI lifespan shutdown cleanup.
|
|
|
|
On an abrupt shutdown (Windows console-close, interpreter teardown racing
|
|
uvicorn) the loop's default executor may already be dead, so an unguarded
|
|
``asyncio.to_thread`` raise here would abort the nested-lifespan unwind and
|
|
surface as "Application shutdown failed". Dependency-injected so it can be
|
|
unit-tested without the heavy backend import graph.
|
|
"""
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import types
|
|
from typing import Callable
|
|
|
|
import structlog
|
|
|
|
logger = structlog.get_logger(__name__)
|
|
|
|
|
|
async def run_lifespan_shutdown(
|
|
terminate_downloads: Callable[[], None],
|
|
clear_compiled_cache: Callable[[], None],
|
|
hw_module: types.ModuleType,
|
|
) -> None:
|
|
"""Run each shutdown step guarded so one failure can't skip the others; never raise."""
|
|
loop = asyncio.get_running_loop()
|
|
# Copy context for parity with asyncio.to_thread. Schedule and await
|
|
# separately so a dead executor (raises at submit) runs inline, while a
|
|
# body exception (raised at await) is logged, not re-run.
|
|
ctx = contextvars.copy_context()
|
|
try:
|
|
future = loop.run_in_executor(None, ctx.run, terminate_downloads)
|
|
except RuntimeError:
|
|
# Executor gone: run inline on the loop thread.
|
|
try:
|
|
ctx.run(terminate_downloads)
|
|
except Exception as exc:
|
|
logger.warning("terminate_downloads (inline) failed at shutdown: %s", exc)
|
|
else:
|
|
try:
|
|
await future
|
|
except Exception as exc:
|
|
logger.warning("terminate_downloads failed at shutdown: %s", exc)
|
|
|
|
try:
|
|
hw_module.DEVICE = None
|
|
except Exception as exc:
|
|
logger.warning("clearing hardware DEVICE failed at shutdown: %s", exc)
|
|
|
|
try:
|
|
clear_compiled_cache()
|
|
except Exception as exc:
|
|
logger.warning("clear_compiled_cache failed at shutdown: %s", exc)
|