1  import signal
 2  from contextlib import contextmanager
 3  
 4  from mlflow.exceptions import MlflowException
 5  from mlflow.protos.databricks_pb2 import NOT_IMPLEMENTED
 6  from mlflow.utils.os import is_windows
 7  
 8  
 9  class MlflowTimeoutError(Exception):
10      pass
11  
12  
13  @contextmanager
14  def run_with_timeout(seconds):
15      """
16      Context manager to runs a block of code with a timeout. If the block of code takes longer
17      than `seconds` to execute, a `TimeoutError` is raised.
18      NB: This function uses Unix signals to implement the timeout, so it is not thread-safe.
19      Also it does not work on non-Unix platforms such as Windows.
20  
21      E.g.
22          ```
23          with run_with_timeout(5):
24              model.predict(data)
25          ```
26      """
27      if is_windows():
28          raise MlflowException(
29              "Timeouts are not implemented yet for non-Unix platforms",
30              error_code=NOT_IMPLEMENTED,
31          )
32  
33      def signal_handler(signum, frame):
34          raise MlflowTimeoutError(f"Operation timed out after {seconds} seconds")
35  
36      signal.signal(signal.SIGALRM, signal_handler)
37      signal.alarm(seconds)
38  
39      try:
40          yield
41      finally:
42          signal.alarm(0)  # Disable the alarm after the operation completes or times out