diff --git a/rag/svr/task_executor.py b/rag/svr/task_executor.py index 320ac130b..247042641 100644 --- a/rag/svr/task_executor.py +++ b/rag/svr/task_executor.py @@ -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)