Add unit tests

Signed-off-by: ashbhandare <[email protected]>

Author: ashbhandare

nemo_run/core/execution/base.py

@@ -166,7 +166,7 @@ def get_launcher_prefix(self) -> Optional[list[str]]:
             return launcher.get_nsys_prefix(profile_dir=self.job_dir)
 
     def get_nsys_entrypoint(self) -> str:
-        return ("nsys","")
+        return ("nsys", "")
 
     def supports_launcher_transform(self) -> bool:
         return False

nemo_run/core/execution/slurm.py

@@ -551,13 +551,13 @@ def get_launcher_prefix(self) -> Optional[list[str]]:
             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="'"
+            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:

test/core/execution/test_base.py

@@ -126,6 +126,11 @@ def test_get_launcher_str(self):
         executor = Executor(launcher="torchrun")
         assert isinstance(executor.get_launcher(), Torchrun)
 
+    def test_get_nsys_entrypoint(self):
+        mock_launcher = Launcher()
+        executor = Executor(launcher=mock_launcher)
+        assert executor.get_nsys_entrypoint() == ("nsys", "")
+
     def test_cleanup(self):
         executor = Executor()
         assert executor.cleanup("handle") is None

test/core/execution/test_slurm.py

@@ -170,10 +170,38 @@ def test_get_launcher_prefix(self):
         launcher_mock = MagicMock()
         launcher_mock.nsys_profile = True
         launcher_mock.get_nsys_prefix.return_value = ["nsys", "profile"]
+        launcher_mock.nsys_gpu_metrics = False
 
         with patch.object(executor, "get_launcher", return_value=launcher_mock):
             assert executor.get_launcher_prefix() == ["nsys", "profile"]
 
+    def test_get_launcher_prefix_with_gpu_metrics(self):
+        """Test the get_launcher_prefix method with nsys_profile when gpu metrics is enabled."""
+        executor = SlurmExecutor(account="test")
+
+        # Test with launcher that has nsys_profile
+        launcher_mock = MagicMock()
+        launcher_mock.nsys_profile = True
+        launcher_mock.get_nsys_prefix.return_value = ["nsys", "profile"]
+        launcher_mock.nsys_gpu_metrics = True
+
+        with patch.object(executor, "get_launcher", return_value=launcher_mock):
+            assert executor.get_launcher_prefix() == ["nsys", "profile", "$GPU_METRICS_FLAG"]
+
+    def test_get_nsys_entrypoint(self):
+        """Test the get_nsys_entrypoint method with nsys_profile."""
+        executor = SlurmExecutor(account="test")
+
+        # Test with launcher that has nsys_profile
+        launcher_mock = MagicMock()
+        launcher_mock.nsys_gpu_metrics = True
+
+        with patch.object(executor, "get_launcher", return_value=launcher_mock):
+            assert executor.get_nsys_entrypoint() == (
+                'bash -c \'GPU_METRICS_FLAG=""; if [ "$SLURM_PROCID" -eq 0 ]; then GPU_METRICS_FLAG="--gpu-metrics-devices=all"; fi; nsys',
+                "'",
+            )
+
     def test_supports_launcher_transform(self):
         """Test the supports_launcher_transform method."""
         executor = SlurmExecutor(account="test")