GPU metrics only on rank 0

Author: Aishwarya Bhandare

nemo_run/core/execution/base.py

@@ -165,6 +165,9 @@ def get_launcher_prefix(self) -> Optional[list[str]]:
             os.makedirs(os.path.join(self.job_dir, launcher.nsys_folder), exist_ok=True)
             return launcher.get_nsys_prefix(profile_dir=self.job_dir)
 
+    def get_nsys_entrypoint(self) -> str:
+        return ("nsys","")
+
     def supports_launcher_transform(self) -> bool:
         return False
 

nemo_run/core/execution/launcher.py

@@ -24,6 +24,7 @@ class Launcher(ConfigurableMixin):
             "--cuda-event-trace=false",
         ]
     )
+    nsys_gpu_metrics: bool = False
 
     def get_nsys_prefix(self, profile_dir: str) -> Optional[list[str]]:
         """Make a command prefix for nsys profiling"""

nemo_run/core/execution/slurm.py

@@ -547,7 +547,18 @@ def assign(
     def get_launcher_prefix(self) -> Optional[list[str]]:
         launcher = self.get_launcher()
         if launcher.nsys_profile:
-            return launcher.get_nsys_prefix(profile_dir=f"/{RUNDIR_NAME}")
+            nsys_prefix = launcher.get_nsys_prefix(profile_dir=f"/{RUNDIR_NAME}")
+            if launcher.nsys_gpu_metrics:
+                nsys_prefix += ["$GPU_METRICS_FLAG"]
+        return nsys_prefix
+    
+    def get_nsys_entrypoint(self) -> str:
+        launcher = self.get_launcher()
+        entrypoint, postfix = "nsys", ""
+        if launcher.nsys_gpu_metrics:
+            entrypoint="bash -c 'GPU_METRICS_FLAG=\"\"; if [ \"$SLURM_PROCID\" -eq 0 ]; then GPU_METRICS_FLAG=\"--gpu-metrics-devices=all\"; fi; nsys"
+            postfix="'"
+        return (entrypoint, postfix)
 
     def supports_launcher_transform(self) -> bool:
         return True if isinstance(self.get_launcher(), SlurmTemplate) else False

nemo_run/run/torchx_backend/packaging.py

@@ -225,8 +225,8 @@ def _get_details_from_script(fn_or_script: Script, serialize_configs: bool):
         nsys_prefix = executor.get_launcher_prefix()
         if nsys_prefix:
             role.args = [role.entrypoint] + role.args
-            role.entrypoint = "nsys"
-            role.args = nsys_prefix + role.args
+            role.entrypoint, nsys_postfix = executor.get_nsys_entrypoint()
+            role.args = nsys_prefix + role.args + [nsys_postfix]
 
     if metadata:
         if USE_WITH_RAY_CLUSTER_KEY in metadata: