Spaces:
Sleeping
Sleeping
File size: 1,723 Bytes
2680a94 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
import json
import os
from threading import Thread
import resource
import signal
import json
import os
import traceback
from loguru import logger
def set_memory_limit(max_memory):
def limit_memory():
resource.setrlimit(resource.RLIMIT_AS, (max_memory, max_memory))
return limit_memory
def timeout_handler(_, __):
raise TimeoutError()
def to_jsonl(dict_data, file_path):
with open(file_path, 'a') as file:
json_line = json.dumps(dict_data)
file.write(json_line + os.linesep)
class PropagatingThread(Thread):
def run(self):
self.exc = None
try:
if hasattr(self, '_Thread__target'):
# Thread uses name mangling prior to Python 3.
self.ret = self._Thread__target(*self._Thread__args, **self._Thread__kwargs)
else:
self.ret = self._target(*self._args, **self._kwargs)
except Exception as e:
self.exc = e
def join(self, timeout=None):
super(PropagatingThread, self).join(timeout)
if self.exc:
raise self.exc
if self.is_alive():
return None
return self.ret
def terminate(self):
self._stop()
def function_with_timeout(func, args, timeout, max_memory=100 * 1024 * 1024):
result_container = []
def wrapper():
# set_memory_limit(max_memory)()
result_container.append(func(*args))
thread = PropagatingThread(target=wrapper)
thread.start()
thread.join(timeout)
if thread.is_alive():
logger.error(f"Timeout Error\n {args[0]} with timeout {timeout}")
thread.terminate()
raise TimeoutError()
else:
return result_container[0]
|