From 9a1c7b7b1d6997a5e608110c73ec22477e0f841a Mon Sep 17 00:00:00 2001
From: ST-XX <15625257+ST-XX@users.noreply.github.com>
Date: Tue, 20 Jan 2026 12:03:32 +0800
Subject: [PATCH 1/2] nvtx

---
 fastdeploy/utils.py | 48 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/fastdeploy/utils.py b/fastdeploy/utils.py
index 2796b931c2f..f664f3f0ec7 100644
--- a/fastdeploy/utils.py
+++ b/fastdeploy/utils.py
@@ -1285,3 +1285,51 @@ def decorator(func):
     register_op = do_nothing
 
 register_custom_python_op = register_op
+
+
+from functools import wraps
+
+
+def nvtx_annotate(message, color="blue"):
+    """NVTX"""
+    import nvtx
+
+    def decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            with nvtx.annotate(message, color=color):
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def nvtx_class_annotate(color="blue"):
+    """
+    类装饰器：为类的所有公有方法自动添加 NVTX 注解
+    message 格式：ClassName.method_name
+    """
+
+    def class_decorator(cls):
+        class_name = cls.__name__
+        # 遍历类的所有属性
+        for attr_name in dir(cls):
+            # 过滤掉私有方法（以 _ 开头）和特殊方法
+            if attr_name.startswith("_"):
+                continue
+            attr = getattr(cls, attr_name)
+            # 检查是否为可调用的方法
+            if callable(attr):
+                # 生成注解消息
+                message = f"{class_name}.{attr_name}"
+
+                # 应用 NVTX 装饰器
+                decorated_method = nvtx_annotate(message, color=color)(attr)
+
+                # 将装饰后的方法设置回类
+                setattr(cls, attr_name, decorated_method)
+
+        return cls
+
+    return class_decorator

From 3d9b1e1f2a42f6296389e272054be3eff8e88f54 Mon Sep 17 00:00:00 2001
From: ST-XX <15625257+ST-XX@users.noreply.github.com>
Date: Tue, 20 Jan 2026 14:22:36 +0800
Subject: [PATCH 2/2] nvtx

---
 fastdeploy/utils.py | 32 +++++++++++++++++++++-----------
 1 file changed, 21 insertions(+), 11 deletions(-)

diff --git a/fastdeploy/utils.py b/fastdeploy/utils.py
index f664f3f0ec7..3fd12555cef 100644
--- a/fastdeploy/utils.py
+++ b/fastdeploy/utils.py
@@ -1286,15 +1286,21 @@ def decorator(func):
 
 register_custom_python_op = register_op
 
-
 from functools import wraps
 
+try:
+    import nvtx
+except ImportError:
+    nvtx = None
+
 
 def nvtx_annotate(message, color="blue"):
-    """NVTX"""
-    import nvtx
+    """A decorator to add NVTX annotations for profiling."""
 
     def decorator(func):
+        if nvtx is None:
+            return func
+
         @wraps(func)
         def wrapper(*args, **kwargs):
             with nvtx.annotate(message, color=color):
@@ -1307,27 +1313,31 @@ def wrapper(*args, **kwargs):
 
 def nvtx_class_annotate(color="blue"):
     """
-    类装饰器：为类的所有公有方法自动添加 NVTX 注解
-    message 格式：ClassName.method_name
+    A class decorator that automatically adds NVTX annotations to all public
+    methods of a class. The annotation message will be in the format
+    'ClassName.method_name'.
     """
 
     def class_decorator(cls):
+        if nvtx is None:
+            return cls
+
         class_name = cls.__name__
-        # 遍历类的所有属性
+        # Iterate over all attributes of the class
         for attr_name in dir(cls):
-            # 过滤掉私有方法（以 _ 开头）和特殊方法
+            # Filter out private (starting with '_') and special methods
             if attr_name.startswith("_"):
                 continue
             attr = getattr(cls, attr_name)
-            # 检查是否为可调用的方法
+            # Check if it is a callable method
             if callable(attr):
-                # 生成注解消息
+                # Generate annotation message
                 message = f"{class_name}.{attr_name}"
 
-                # 应用 NVTX 装饰器
+                # Apply NVTX decorator
                 decorated_method = nvtx_annotate(message, color=color)(attr)
 
-                # 将装饰后的方法设置回类
+                # Set the decorated method back to the class
                 setattr(cls, attr_name, decorated_method)
 
         return cls