Unverified Commit f6701bc6 authored by Gary Wang's avatar Gary Wang Committed by GitHub
Browse files

Reset alarm signal when the function is ended (#29706)

Fixes #29690
parent e644b600
......@@ -592,8 +592,9 @@ def resolve_trust_remote_code(trust_remote_code, model_name, has_local_code, has
if has_local_code:
trust_remote_code = False
elif has_remote_code and TIME_OUT_REMOTE_CODE > 0:
prev_sig_handler = None
try:
signal.signal(signal.SIGALRM, _raise_timeout_error)
prev_sig_handler = signal.signal(signal.SIGALRM, _raise_timeout_error)
signal.alarm(TIME_OUT_REMOTE_CODE)
while trust_remote_code is None:
answer = input(
......@@ -614,6 +615,10 @@ def resolve_trust_remote_code(trust_remote_code, model_name, has_local_code, has
f"load the model. You can inspect the repository content at https://hf.co/{model_name}.\n"
f"Please pass the argument `trust_remote_code=True` to allow custom code to be run."
)
finally:
if prev_sig_handler is not None:
signal.signal(signal.SIGALRM, prev_sig_handler)
signal.alarm(0)
elif has_remote_code:
# For the CI which puts the timeout at 0
_raise_timeout_error(None, None)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment