2222from .single import PROGRESS_MIN_TIME
2323from .utils import (
2424 StrPath , TestName ,
25- format_duration , print_warning , count , plural )
25+ format_duration , print_warning , count , plural , get_process_memory_usage )
2626from .worker import create_worker_process , USE_PROCESS_GROUP
2727
2828if MS_WINDOWS :
@@ -452,6 +452,12 @@ def wait_stopped(self, start_time: float) -> None:
452452 print_warning (f"Failed to join { self } in { format_duration (dt )} " )
453453 break
454454
455+ def get_mem_usage (self ):
456+ popen = self ._popen
457+ if popen is None :
458+ return
459+ return get_process_memory_usage (popen .pid )
460+
455461
456462def get_running (workers : list [WorkerThread ]) -> str | None :
457463 running : list [str ] = []
@@ -473,6 +479,7 @@ def __init__(self, num_workers: int, runtests: RunTests,
473479 logger : Logger , results : TestResults ) -> None :
474480 self .num_workers = num_workers
475481 self .runtests = runtests
482+ self .logger = logger
476483 self .log = logger .log
477484 self .display_progress = logger .display_progress
478485 self .results : TestResults = results
@@ -598,9 +605,21 @@ def _process_result(self, item: QueueOutput) -> TestResult:
598605
599606 return result
600607
608+ def get_mem_usage (self ):
609+ usage = 0
610+ main_mem = get_process_memory_usage (os .getpid ())
611+ if main_mem :
612+ usage += main_mem
613+ for worker in self .workers :
614+ worker_mem = worker .get_mem_usage ()
615+ if worker_mem :
616+ usage += worker_mem
617+ return usage
618+
601619 def run (self ) -> None :
602620 fail_fast = self .runtests .fail_fast
603621 fail_env_changed = self .runtests .fail_env_changed
622+ self .logger .get_mem_usage = self .get_mem_usage
604623
605624 self .start_workers ()
606625
@@ -625,3 +644,4 @@ def run(self) -> None:
625644 # worker when we exit this function
626645 self .pending .stop ()
627646 self .stop_workers ()
647+ self .logger .get_mem_usage = None
0 commit comments