Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow add middleware after app has started #16758

Open
wants to merge 1 commit into
base: dev
Choose a base branch
from

Conversation

w-e-w
Copy link
Collaborator

@w-e-w w-e-w commented Dec 28, 2024

allow add middleware after app has started

this should completely fix RuntimeError("Cannot add middleware after an application has started")
which can occur due to a race condition

explanation of race condition

By design Starlette / FastAPI middleware can only be added before the the server has started, technically it can still be added until the server receive a first request, at which point it builds the app.middleware_stack
after this point you normally cannot add middleware

BUT we do add and remove middlewares after the app is launched

app, local_url, share_url = shared.demo.launch(

app.user_middleware = [x for x in app.user_middleware if x.cls.__name__ != 'CORSMiddleware']
initialize_util.setup_middleware(app)
progress.setup_progress_api(app)
ui.setup_ui_api(app)
if launch_api:
create_api(app)

def setup_middleware(app):
from starlette.middleware.gzip import GZipMiddleware
app.middleware_stack = None # reset current middleware to allow modifying user provided list
app.add_middleware(GZipMiddleware, minimum_size=1000)
configure_cors_middleware(app)
app.build_middleware_stack() # rebuild middleware stack on-the-fly

in this section of code we first remove the gradio default CORSMiddleware then added our own CORSMiddleware with different config along with GZipMiddleware

and after this if --api is enabled we also add some more middlewares

def api_middleware(app: FastAPI):
rich_available = False
try:
if os.environ.get('WEBUI_RICH_EXCEPTIONS', None) is not None:
import anyio # importing just so it can be placed on silent list
import starlette # importing just so it can be placed on silent list
from rich.console import Console
console = Console()
rich_available = True
except Exception:
pass
@app.middleware("http")
async def log_and_time(req: Request, call_next):
ts = time.time()
res: Response = await call_next(req)
duration = str(round(time.time() - ts, 4))
res.headers["X-Process-Time"] = duration
endpoint = req.scope.get('path', 'err')
if shared.cmd_opts.api_log and endpoint.startswith('/sdapi'):
print('API {t} {code} {prot}/{ver} {method} {endpoint} {cli} {duration}'.format(
t=datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S.%f"),
code=res.status_code,
ver=req.scope.get('http_version', '0.0'),
cli=req.scope.get('client', ('0:0.0.0', 0))[0],
prot=req.scope.get('scheme', 'err'),
method=req.scope.get('method', 'err'),
endpoint=endpoint,
duration=duration,
))
return res
def handle_exception(request: Request, e: Exception):
err = {
"error": type(e).__name__,
"detail": vars(e).get('detail', ''),
"body": vars(e).get('body', ''),
"errors": str(e),
}
if not isinstance(e, HTTPException): # do not print backtrace on known httpexceptions
message = f"API error: {request.method}: {request.url} {err}"
if rich_available:
print(message)
console.print_exception(show_locals=True, max_frames=2, extra_lines=1, suppress=[anyio, starlette], word_wrap=False, width=min([console.width, 200]))
else:
errors.report(message, exc_info=True)
return JSONResponse(status_code=vars(e).get('status_code', 500), content=jsonable_encoder(err))
@app.middleware("http")
async def exception_handling(request: Request, call_next):
try:
return await call_next(request)
except Exception as e:
return handle_exception(request, e)
@app.exception_handler(Exception)
async def fastapi_exception_handler(request: Request, e: Exception):
return handle_exception(request, e)
@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, e: HTTPException):
return handle_exception(request, e)

we are able to add new middlewares because in setup_middleware() we set app.middleware_stack to None, which would make Starlette / FastAPI rebuild the app.middleware_stack on the next recived request

In most cases this works fine, but if one is unlucky enough it is totally possible for the server to receive a request after we set app.middleware_stack to None and between the last middleware is added
if this were to happen it will trigger a runtime error which would then crash the server


reproducing the race condition

key point

if ther is a request between app.middleware_stack = None but and add_middleware() then the issue will be trigger

For testing purposes this issue can be manually triggered by placeing a breakpoin in setup_middleware() after app.middleware_stack = None but before any app.add_middleware
for example on

app.add_middleware(GZipMiddleware, minimum_size=1000)

When the breakpoint is hit, make a request the server (visit the webui webpage or make any request on any endpoint
or perform any request on any API route)
after request continued execution

you should see RuntimeError("Cannot add middleware after an application has started") been trigger

this is because even though we have set app.middleware_stack = None if a request has been made to the server it will rebuild middleware_stack, makeing middleware_stack no longer None, triggeing the error

Is this an issue that needs addressing?

Yes
well this doesn't happen often, it does happen frequent enough that we have seen multiple report of this issue

also remember seeing some reports of the server dying soon after launch, this may be the cause of those issues,
when runtime error is triggered the server only shutsdown after the model is loaded, resulting in a delay between the runtime error and the actual shutdown

depending on the storage speed it is totally probable to have enough time to open the webpage as if everything is normal only to have the server shutdown soon after causing confusion

browser auto launch after app start, would significantly increase the chance of this happening
as it essentially is sending lots of requests to the serve at around the same time as we add new middleare

I suspect this would cause this issue to be triggered more often on slower systems


The Fix

the fix I came up with is to patch Starlette.add_middleware, adding some logic that ensures the middleware_stack is None when we add new middlewares

the resulting of this is that I believe it is now impossible for it to trigger a RuntimeError due to this issue
this even makes it possible to add new middleware at any time

I made a test extention to test adding of middleware after an application has started
https://gist.github.com/w-e-w/3b98e4fe30b3bf75c1f632bdb1ec704e

Remarks

while setting middleware_stack to None makeing it rebuild the stack is a but unorthodox so could be considered dangerous and could break in the future
BUT, but considering we have been doing this operation for a long time and it has worked I think it's good enough
The only difference that this PR does is add better logic in performing this operation

maybe it will cause trouble when we upgrade gradio, but I don't know if that would ever come

Checklist:

this should completely fix "Cannot add middleware after an application has started" which can occur due to a race condition
@w-e-w w-e-w changed the title allow add middleware after app has started allow add middleware after app has started, fix RuntimeError("Cannot add middleware after an application has started") Dec 28, 2024
@w-e-w w-e-w changed the title allow add middleware after app has started, fix RuntimeError("Cannot add middleware after an application has started") Allow add middleware after app has started, fix RuntimeError("Cannot add middleware after an application has started") Dec 28, 2024
@w-e-w w-e-w changed the title Allow add middleware after app has started, fix RuntimeError("Cannot add middleware after an application has started") Allow add middleware after app has started Dec 28, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant