Fix nvtx name for vllm internal executor and add nvtx for vllm worker

Signed-off-by: Guyue Huang <guyueh@nvidia.com>

Author: guyueh1

nemo_rl/models/generation/vllm.py

@@ -53,6 +53,7 @@
 )
 from nemo_rl.models.huggingface.common import ModelFlag
 from nemo_rl.models.policy.utils import is_vllm_v1_engine_enabled
+from nemo_rl.utils.nsys import wrap_with_nvtx_name
 
 
 class VllmSpecificArgs(TypedDict):
@@ -448,6 +449,7 @@ def _build_sampling_params(
             include_stop_str_in_output=True,
         )
 
+    @wrap_with_nvtx_name("vllm_genertion_worker/generate")
     def generate(
         self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False
     ) -> BatchedDataDict[GenerationOutputSpec]:
@@ -811,6 +813,7 @@ async def process_single_sample(sample_idx):
                 await asyncio.gather(*sample_tasks, return_exceptions=True)
                 raise e
 
+    @wrap_with_nvtx_name("vllm_genertion_worker/generate_text")
     def generate_text(
         self, data: BatchedDataDict[GenerationDatumSpec], greedy: bool = False
     ) -> BatchedDataDict[GenerationOutputSpec]:
@@ -1045,6 +1048,7 @@ async def prepare_refit_info_async(self, state_dict_info: dict[str, Any]) -> Non
         """Async version of prepare_refit_info."""
         await self.llm.collective_rpc("prepare_refit_info", args=(state_dict_info,))
 
+    @wrap_with_nvtx_name("vllm_genertion_worker/update_weights_from_ipc_handles")
     def update_weights_from_ipc_handles(self, ipc_handles: dict[str, Any]) -> bool:
         """Update weights from IPC handles by delegating to the vLLM Worker implementation.
 
@@ -1156,6 +1160,7 @@ async def update_weights_from_ipc_handles_async(
             traceback.print_exc()
             return False
 
+    @wrap_with_nvtx_name("vllm_genertion_worker/update_weights_from_collective")
     def update_weights_from_collective(self) -> bool:
         """Update the model weights from collective communication."""
         try:

nemo_rl/models/generation/vllm_backend.py

@@ -68,7 +68,9 @@ def prepare_refit_info(
         """
         self.state_dict_info = state_dict_info  # pyrefly: ignore[implicitly-defined-attribute]  This class does not define __init__ so assignments like this should be ignored
 
-    @wrap_with_nvtx_name("update_weights_from_global_ipc_handles")
+    @wrap_with_nvtx_name(
+        "vllm_internal_worker_extension/update_weights_from_global_ipc_handles"
+    )
     def update_weights_from_global_ipc_handles(self, global_device_ipc_handles):
         """Update weights from global IPC handles.
 
@@ -82,7 +84,9 @@ def update_weights_from_global_ipc_handles(self, global_device_ipc_handles):
         local_device_ipc_handles = global_device_ipc_handles[device_uuid]
         return self.update_weights_from_local_ipc_handles(local_device_ipc_handles)
 
-    @wrap_with_nvtx_name("update_weights_from_local_ipc_handles")
+    @wrap_with_nvtx_name(
+        "vllm_internal_worker_extension/update_weights_from_local_ipc_handles"
+    )
     def update_weights_from_local_ipc_handles(self, local_device_ipc_handles):
         """Update weights from local IPC handles.
 
@@ -159,7 +163,9 @@ def update_weights_from_local_ipc_handles(self, local_device_ipc_handles):
             )
             return False
 
-    @wrap_with_nvtx_name("update_weights_from_collective")
+    @wrap_with_nvtx_name(
+        "vllm_internal_worker_extension/update_weights_from_collective"
+    )
     def update_weights_from_collective(self) -> bool:
         """Update the model weights from collective communication."""
         assert self.state_dict_info is not None, (