From 78c9b8e3c9621bef5e79c8f28972c8479fef6358 Mon Sep 17 00:00:00 2001
From: Jacek Bieniusiewicz <jbieniusiewi@nvidia.com>
Date: Wed, 22 Jan 2025 13:54:21 +0100
Subject: [PATCH] Lazy loading of cupti_module

---
 src/nvidia_resiliency_ext/straggler/cupti.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/nvidia_resiliency_ext/straggler/cupti.py b/src/nvidia_resiliency_ext/straggler/cupti.py
index b268dd9..9b79a8a 100644
--- 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,