Add tracing feature

ceci/stage.py

@@ -14,6 +14,7 @@
 from . import errors
 from .monitor import MemoryMonitor
 from .config import StageParameter, StageConfig, cast_to_streamable
+from .utils import activate_tracing
 
 SERIAL = "serial"
 MPI_PARALLEL = "mpi"
@@ -610,6 +611,12 @@ def parse_command_line(cls, cmd=None):
             help="Report memory use. Argument gives interval in seconds between reports",
         )
 
+        parser.add_argument(
+            "--trace",
+            action="store_true",
+            help="Enable sending a signal to the process that prints a trace wherever it is",
+        )
+
         # Error message we will return if --mpi used on a non-supported
         # stage.
         mpi_err = (
@@ -668,6 +675,9 @@ def execute(cls, args, comm=None):
         if args.memmon:  # pragma: no cover
             monitor = MemoryMonitor.start_in_thread(interval=args.memmon)
 
+        if args.trace:
+            activate_tracing(stage._rank)
+
         # Now we try to see if the validation step has been changed,
         # if it has then we will run the validation step, and raise any errors
         try:

ceci/utils.py

@@ -3,6 +3,8 @@
 from contextlib import contextmanager
 import sys
 import os
+import signal
+import faulthandler
 
 
 def add_python_path(path, start):
@@ -107,3 +109,9 @@ def embolden(text):
         Emboldened text
     """
     return "\033[1m" + text + "\033[0m"
+
+
+def activate_tracing(rank):
+    pid = os.getpid()
+    print(f"Activating trace mode for process {rank}. Trigger trace using command: kill -s {signal.SIGUSR1} {pid}")
+    faulthandler.register(signal.SIGUSR1)