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 ...@@ -592,8 +592,9 @@ def resolve_trust_remote_code(trust_remote_code, model_name, has_local_code, has
if has_local_code: if has_local_code:
trust_remote_code = False trust_remote_code = False
elif has_remote_code and TIME_OUT_REMOTE_CODE > 0: elif has_remote_code and TIME_OUT_REMOTE_CODE > 0:
prev_sig_handler = None
try: try:
signal.signal(signal.SIGALRM, _raise_timeout_error) prev_sig_handler = signal.signal(signal.SIGALRM, _raise_timeout_error)
signal.alarm(TIME_OUT_REMOTE_CODE) signal.alarm(TIME_OUT_REMOTE_CODE)
while trust_remote_code is None: while trust_remote_code is None:
answer = input( answer = input(
...@@ -614,6 +615,10 @@ def resolve_trust_remote_code(trust_remote_code, model_name, has_local_code, has ...@@ -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"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." 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: elif has_remote_code:
# For the CI which puts the timeout at 0 # For the CI which puts the timeout at 0
_raise_timeout_error(None, None) _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