|
|
|
@@ -40,7 +40,6 @@ from io import BytesIO |
|
|
|
from multiprocessing.context import TimeoutError |
|
|
|
from timeit import default_timer as timer |
|
|
|
import tracemalloc |
|
|
|
import resource |
|
|
|
import signal |
|
|
|
import trio |
|
|
|
import exceptiongroup |
|
|
|
@@ -117,7 +116,13 @@ def start_tracemalloc_and_snapshot(signum, frame): |
|
|
|
snapshot = tracemalloc.take_snapshot() |
|
|
|
snapshot.dump(snapshot_file) |
|
|
|
current, peak = tracemalloc.get_traced_memory() |
|
|
|
max_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
|
|
|
if sys.platform == "win32": |
|
|
|
import psutil |
|
|
|
process = psutil.Process() |
|
|
|
max_rss = process.memory_info().rss / 1024 |
|
|
|
else: |
|
|
|
import resource |
|
|
|
max_rss = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
|
|
|
logging.info(f"taken snapshot {snapshot_file}. max RSS={max_rss / 1000:.2f} MB, current memory usage: {current / 10**6:.2f} MB, Peak memory usage: {peak / 10**6:.2f} MB") |
|
|
|
|
|
|
|
# SIGUSR2 handler: stop tracemalloc |
|
|
|
@@ -644,8 +649,9 @@ async def main(): |
|
|
|
logging.info(f'TaskExecutor: RAGFlow version: {get_ragflow_version()}') |
|
|
|
settings.init_settings() |
|
|
|
print_rag_settings() |
|
|
|
signal.signal(signal.SIGUSR1, start_tracemalloc_and_snapshot) |
|
|
|
signal.signal(signal.SIGUSR2, stop_tracemalloc) |
|
|
|
if sys.platform != "win32": |
|
|
|
signal.signal(signal.SIGUSR1, start_tracemalloc_and_snapshot) |
|
|
|
signal.signal(signal.SIGUSR2, stop_tracemalloc) |
|
|
|
TRACE_MALLOC_ENABLED = int(os.environ.get('TRACE_MALLOC_ENABLED', "0")) |
|
|
|
if TRACE_MALLOC_ENABLED: |
|
|
|
start_tracemalloc_and_snapshot(None, None) |