Lazy loading of cupti_module

NVIDIA · Jan 22, 2025 · 78c9b8e · 78c9b8e
1 parent 98ebc42
commit 78c9b8e
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/nvidia_resiliency_ext/straggler/cupti.py b/src/nvidia_resiliency_ext/straggler/cupti.py
@@ -15,8 +15,6 @@
 
 import threading
 
-from . import cupti_module  # type: ignore
-
 
 class CuptiManager:
     """Provide thread safe access to the CUPTI extension.
@@ -32,6 +30,9 @@ def __init__(self, bufferSize=1_000_000, numBuffers=8, statsMaxLenPerKernel=4096
             statsMaxLenPerKernel (int, optional): Max number of timing entries per kernel.
                 (when this limit is rached, oldest timing entries are discarded). Defaults to 4096.
         """
+        # lazy load to avoid circular import
+        from . import cupti_module  # type: ignore
+
         self.cupti_ext = cupti_module.CuptiProfiler(
             bufferSize=bufferSize,
             numBuffers=numBuffers,